ImportClassHelper (quick fixes) refactoring

idea/src/org/jetbrains/jet/plugin/JetPluginUtil.java

@@ -4,14 +4,12 @@ import com.google.common.collect.Lists;
 import com.intellij.openapi.project.Project;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
+import org.jetbrains.jet.lang.descriptors.ModuleDescriptor;
 import org.jetbrains.jet.lang.descriptors.NamespaceDescriptor;
 import org.jetbrains.jet.lang.resolve.java.JavaClassDescriptor;
 import org.jetbrains.jet.lang.resolve.java.JavaDescriptorResolver;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
-import org.jetbrains.jet.lang.types.DeferredType;
-import org.jetbrains.jet.lang.types.ErrorUtils;
-import org.jetbrains.jet.lang.types.JetStandardLibrary;
-import org.jetbrains.jet.lang.types.JetType;
+import org.jetbrains.jet.lang.types.*;
 
 import java.util.LinkedList;
 
@@ -41,7 +39,7 @@ public class JetPluginUtil {
         DeclarationDescriptor declarationDescriptor = type.getConstructor().getDeclarationDescriptor();
 
         LinkedList<String> fullName = Lists.newLinkedList();
-        while (declarationDescriptor != null) {
+        while (declarationDescriptor != null && !(declarationDescriptor instanceof ModuleDescriptor)) {
             fullName.addFirst(declarationDescriptor.getName());
             declarationDescriptor = declarationDescriptor.getContainingDeclaration();
         }
@@ -53,6 +51,11 @@ public class JetPluginUtil {
     }
 
     public static boolean checkTypeIsStandard(JetType type, Project project) {
+        if (JetStandardClasses.isAny(type) || JetStandardClasses.isNothingOrNullableNothing(type) || JetStandardClasses.isUnit(type) ||
+            JetStandardClasses.isTupleType(type) || JetStandardClasses.isFunctionType(type)) {
+            return true;
+        }
+
         LinkedList<String> fullName = computeTypeFullNameList(type);
         if (fullName.size() == 3 && fullName.getFirst().equals("java") && fullName.get(1).equals("lang")) {
             return true;

idea/src/org/jetbrains/jet/plugin/actions/JetAddImportAction.java

@@ -13,8 +13,10 @@ import com.intellij.openapi.ui.popup.PopupStep;
 import com.intellij.openapi.ui.popup.util.BaseListPopupStep;
 import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.util.PlatformIcons;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.plugin.quickfix.ImportClassHelper;
 
 import javax.swing.*;
@@ -138,9 +140,11 @@ public class JetAddImportAction implements QuestionAction {
                     public void run() {
                         // TODO: See {@link com.intellij.codeInsight.daemon.impl.actions.AddImportAction#_addImport} for more ideas.
                         // TODO: Optimize imports
+                        PsiFile file = element.getContainingFile();
+                        if (!(file instanceof JetFile)) return;
                         ImportClassHelper.addImportDirective(
                                 selectedImport,
-                                ImportClassHelper.findOuterNamespace(element)
+                                (JetFile)file
                         );
                     }
                 });

idea/src/org/jetbrains/jet/plugin/quickfix/AddReturnTypeFix.java

@@ -11,6 +11,8 @@ import org.jetbrains.jet.lang.diagnostics.DiagnosticParameters;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticWithParameters;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticWithPsiElement;
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.types.ErrorUtils;
+import org.jetbrains.jet.lang.types.JetStandardClasses;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.plugin.JetBundle;
 
@@ -37,8 +39,14 @@ public class AddReturnTypeFix extends JetIntentionAction<JetNamedDeclaration> {
         return JetBundle.message("add.return.type");
     }
 
+    @Override
+    public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+        return super.isAvailable(project, editor, file) && !ErrorUtils.isErrorType(type);
+    }
+
     @Override
     public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+        if (!(file instanceof JetFile)) return;
         PsiElement newElement;
         if (element instanceof JetProperty) {
             newElement = addPropertyType(project, (JetProperty) element, type);
@@ -47,7 +55,8 @@ public class AddReturnTypeFix extends JetIntentionAction<JetNamedDeclaration> {
             assert element instanceof JetFunction;
             newElement = addFunctionType(project, (JetFunction) element, type);
         }
-        ImportClassHelper.perform(type, element, newElement);
+        ImportClassHelper.addImportDirectiveIfNeeded(type, (JetFile)file);
+        element.replace(newElement);
     }
 
     @Override

idea/src/org/jetbrains/jet/plugin/quickfix/ChangeAccessorTypeFix.java

@@ -10,10 +10,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticParameters;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticWithParameters;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticWithPsiElement;
-import org.jetbrains.jet.lang.psi.JetParameter;
-import org.jetbrains.jet.lang.psi.JetPropertyAccessor;
-import org.jetbrains.jet.lang.psi.JetPsiFactory;
-import org.jetbrains.jet.lang.psi.JetTypeReference;
+import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.plugin.JetBundle;
 
@@ -57,7 +54,7 @@ public class ChangeAccessorTypeFix extends JetIntentionAction<JetPropertyAccesso
             assert typeReference != null;
             CodeEditUtil.replaceChild(parameter.getNode(), typeReference.getNode(), newTypeReference.getNode());
         }
-        ImportClassHelper.perform(type, element, newElement);
+        element.replace(newElement);
     }
     
     public static JetIntentionActionFactory<JetPropertyAccessor> createFactory() {

idea/src/org/jetbrains/jet/plugin/quickfix/ImportClassHelper.java

@@ -2,11 +2,12 @@ package org.jetbrains.jet.plugin.quickfix;
 
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.JetDeclaration;
-import org.jetbrains.jet.lang.psi.JetImportDirective;
 import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.psi.JetImportDirective;
 import org.jetbrains.jet.lang.psi.JetPsiFactory;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.java.AnalyzerFacade;
 import org.jetbrains.jet.lang.types.ErrorUtils;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.plugin.JetPluginUtil;
@@ -17,48 +18,34 @@ import java.util.List;
  * @author svtk
  */
 public class ImportClassHelper {
-    public static void perform(@NotNull JetType type, @NotNull PsiElement elementToReplace, @NotNull PsiElement newElement) {
-        if (JetPluginUtil.checkTypeIsStandard(type, elementToReplace.getProject()) || ErrorUtils.isError(type.getMemberScope().getContainingDeclaration())) {
-            elementToReplace.replace(newElement);
+    /**
+     * Add import directive corresponding to a type to file when it is needed.
+     *
+     * @param type type to import
+     * @param file file where import directive should be added
+     */
+    public static void addImportDirectiveIfNeeded(@NotNull JetType type, @NotNull JetFile file) {
+        if (JetPluginUtil.checkTypeIsStandard(type, file.getProject()) || ErrorUtils.isErrorType(type)) {
             return;
         }
-        perform(JetPluginUtil.computeTypeFullName(type), elementToReplace, newElement);
-    }
-    
-    public static void perform(@NotNull String typeFullName, @NotNull JetFile namespace) {
-        perform(typeFullName, namespace, null, null);
-    }
-
-    /**
-     * Get the outer namespace PSI element for given element in the tree.
-     *
-     * @param element Some element in the tree.
-     * @return A namespace element in the tree.
-     */
-    public static JetFile findOuterNamespace(@NotNull PsiElement element) {
-        PsiElement parent = element;
-        while (!(parent instanceof JetFile)) {
-            parent = parent.getParent();
-            assert parent != null;
+        BindingContext bindingContext = AnalyzerFacade.analyzeFileWithCache(file, AnalyzerFacade.SINGLE_DECLARATION_PROVIDER);
+        PsiElement element = bindingContext.get(BindingContext.DESCRIPTOR_TO_DECLARATION, type.getMemberScope().getContainingDeclaration());
+        if (element != null && element.getContainingFile() == file) { //declaration is in the same file, so no import is needed
+            return;
         }
-
-        return (JetFile) parent;
-    }
-
-    public static void perform(@NotNull String typeFullName, @NotNull PsiElement elementToReplace, @NotNull PsiElement newElement) {
-        perform(typeFullName, findOuterNamespace(elementToReplace), elementToReplace, newElement);
+        addImportDirective(JetPluginUtil.computeTypeFullName(type), file);
     }
 
     /**
      * Add import directive into the PSI tree for the given namespace.
      *
      * @param importString full name of the import. Can contain .* if necessary.
-     * @param namespace Namespace where directive should be added.
+     * @param file File where directive should be added.
      */
-    public static void addImportDirective(@NotNull String importString, @NotNull JetFile namespace) {
-        List<JetImportDirective> importDirectives = namespace.getImportDirectives();
+    public static void addImportDirective(@NotNull String importString, @NotNull JetFile file) {
+        List<JetImportDirective> importDirectives = file.getImportDirectives();
 
-        JetImportDirective newDirective = JetPsiFactory.createImportDirective(namespace.getProject(), importString);
+        JetImportDirective newDirective = JetPsiFactory.createImportDirective(file.getProject(), importString);
 
         String lineSeparator = System.getProperty("line.separator");
         if (!importDirectives.isEmpty()) {
@@ -72,22 +59,14 @@ public class ImportClassHelper {
 
             JetImportDirective lastDirective = importDirectives.get(importDirectives.size() - 1);
             lastDirective.getParent().addAfter(newDirective, lastDirective);
-            lastDirective.getParent().addAfter(JetPsiFactory.createWhiteSpace(namespace.getProject(), lineSeparator), lastDirective);
+            lastDirective.getParent().addAfter(JetPsiFactory.createWhiteSpace(file.getProject(), lineSeparator), lastDirective);
         }
         else {
-            List<JetDeclaration> declarations = namespace.getDeclarations();
+            List<JetDeclaration> declarations = file.getDeclarations();
             assert !declarations.isEmpty();
             JetDeclaration firstDeclaration = declarations.iterator().next();
             firstDeclaration.getParent().addBefore(newDirective, firstDeclaration);
-            firstDeclaration.getParent().addBefore(JetPsiFactory.createWhiteSpace(namespace.getProject(), lineSeparator + lineSeparator), firstDeclaration);
-        }
-    }
-
-    public static void perform(@NotNull String typeFullName, @NotNull JetFile namespace, @Nullable PsiElement elementToReplace, @Nullable PsiElement newElement) {
-        addImportDirective(typeFullName, namespace);
-
-        if (elementToReplace != null && newElement != null && elementToReplace != newElement) {
-            elementToReplace.replace(newElement);
+            firstDeclaration.getParent().addBefore(JetPsiFactory.createWhiteSpace(file.getProject(), lineSeparator + lineSeparator), firstDeclaration);
         }
     }
 }

idea/src/org/jetbrains/jet/plugin/quickfix/RemovePartsFromPropertyFix.java

@@ -78,6 +78,7 @@ public class RemovePartsFromPropertyFix extends JetIntentionAction<JetProperty>
 
     @Override
     public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+        if (!(file instanceof JetFile)) return;
         JetProperty newElement = (JetProperty) element.copy();
         JetPropertyAccessor getter = newElement.getGetter();
         if (removeGetter && getter != null) {
@@ -102,10 +103,9 @@ public class RemovePartsFromPropertyFix extends JetIntentionAction<JetProperty>
             }
         }
         if (needImport) {
-            ImportClassHelper.perform(type, element, newElement);
-        } else {
-            element.replace(newElement);
+            ImportClassHelper.addImportDirectiveIfNeeded(type, (JetFile)file);
         }
+        element.replace(newElement);
     }
 
     public static JetIntentionActionFactory<JetProperty> createFactory() {

idea/testData/quickfix/typeAddition/beforeNoAddErrorType.kt

@@ -0,0 +1,5 @@
+// "Add return type declaration" "false"
+
+class A() {
+    public val <caret>t = foo()
+}

idea/testData/quickfix/typeImports/afterImportFromAnotherFile.kt

@@ -0,0 +1,9 @@
+// "Add return type declaration" "true"
+
+package a
+
+import b.B
+
+class A() {
+    public val <caret>a : B = b.foo()
+}

idea/testData/quickfix/typeImports/afterNoImportFromTheSameFile.kt

@@ -0,0 +1,7 @@
+// "Add return type declaration" "true"
+
+class A() {}
+
+class B() {
+    public val <caret>a : A = A()
+}

idea/testData/quickfix/typeImports/beforeImportFromAnotherFile.Data.b.kt

@@ -0,0 +1,5 @@
+package b
+
+class B() {}
+
+fun foo() = B()

idea/testData/quickfix/typeImports/beforeImportFromAnotherFile.Main.kt

@@ -0,0 +1,7 @@
+// "Add return type declaration" "true"
+
+package a
+
+class A() {
+    public val <caret>a = b.foo()
+}

idea/testData/quickfix/typeImports/beforeNoImportFromTheSameFile.kt

@@ -0,0 +1,7 @@
+// "Add return type declaration" "true"
+
+class A() {}
+
+class B() {
+    public val <caret>a = A()
+}

idea/tests/org/jetbrains/jet/plugin/quickfix/JetPsiCheckerMultifileTest.java

@@ -147,7 +147,7 @@ public class JetPsiCheckerMultifileTest extends DaemonAnalyzerTestCase {
             }
         });
 
-        final ArrayList<File> fileResult = new ArrayList<File>(allTestFiles);
+        final ArrayList<File> fileResult = new ArrayList<File>(mainFiles);
         fileResult.addAll(dataFiles);
 
         return fileResult;