ChangeFunctionReturnTypeFix quickfix

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.java

@@ -75,6 +75,12 @@ public class JetPsiFactory {
         return Pair.create(property.findElementAt(5), property.findElementAt(7));
     }
 
+    //the pair contains the first and the last elements of a range
+    public static Pair<PsiElement, PsiElement> createTypeWhiteSpaceAndColon(Project project, String type) {
+        JetProperty property = createProperty(project, "val x: " + type);
+        return Pair.create(property.findElementAt(5), (PsiElement) property.getTypeRef());
+    }
+
     public static ASTNode createColonNode(Project project) {
         JetProperty property = createProperty(project, "val x: Int");
         return property.getNode().findChildByType(JetTokens.COLON);

idea/src/org/jetbrains/jet/plugin/JetBundle.properties

@@ -55,6 +55,8 @@ rename.parameter.to.match.overridden.method=Rename parameter to match overridden
 rename.family=Rename
 add.semicolon.after.invocation=Add semicolon after invocation of ''{0}''
 add.semicolon.family=Add Semicolon
+change.function.return.type=Change ''{0}'' function return type to ''{1}''
+remove.function.return.type=Remove explicitly specified return type in ''{0}'' function
 change.return.type.to.match.overridden.method=Change return type to ''{0}''
 remove.return.type.to.match.overridden.method=Remove explicitly specified return type to match overridden method
 change.property.type.to.match.overridden.property=Change property type to ''{0}''

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

@@ -0,0 +1,134 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.plugin.quickfix;
+
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Pair;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiErrorElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.util.IncorrectOperationException;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
+import org.jetbrains.jet.lang.descriptors.SimpleFunctionDescriptor;
+import org.jetbrains.jet.lang.diagnostics.Diagnostic;
+import org.jetbrains.jet.lang.diagnostics.DiagnosticWithParameters3;
+import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.BindingContextUtils;
+import org.jetbrains.jet.lang.resolve.DescriptorResolver;
+import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
+import org.jetbrains.jet.lang.resolve.name.Name;
+import org.jetbrains.jet.lang.types.JetType;
+import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
+import org.jetbrains.jet.plugin.JetBundle;
+import org.jetbrains.jet.plugin.intentions.SpecifyTypeExplicitlyAction;
+import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache;
+
+public class ChangeFunctionReturnTypeFix extends JetIntentionAction<JetFunction> {
+    private final JetType type;
+
+    public ChangeFunctionReturnTypeFix(@NotNull JetFunction element, @NotNull JetType type) {
+        super(element);
+        this.type = type;
+    }
+
+    @NotNull
+    @Override
+    public String getText() {
+        String functionName = element.getName();
+        BindingContext context = AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) element.getContainingFile()).getBindingContext();
+        SimpleFunctionDescriptor descriptor = context.get(BindingContext.FUNCTION, element);
+        if (descriptor != null) {
+            functionName = descriptor.getContainingDeclaration().getName() + "." + functionName;
+        }
+
+        if (KotlinBuiltIns.getInstance().isUnit(type) && element.hasBlockBody()) {
+            return JetBundle.message("remove.function.return.type", functionName);
+        }
+        return JetBundle.message("change.function.return.type", functionName, type.toString());
+    }
+
+    @NotNull
+    @Override
+    public String getFamilyName() {
+        return JetBundle.message("change.type.family");
+    }
+
+    @Override
+    public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+        SpecifyTypeExplicitlyAction.removeTypeAnnotation(element);
+        if (!(KotlinBuiltIns.getInstance().isUnit(type) && element.hasBlockBody())) {
+            PsiElement elementToPrecedeType = element.getValueParameterList();
+            if (elementToPrecedeType == null) elementToPrecedeType = element.getNameIdentifier();
+            assert elementToPrecedeType != null : "Return type of function without name can't mismatch anything";
+            if (elementToPrecedeType.getNextSibling() instanceof PsiErrorElement) {
+                // if a function doesn't have a value parameter list, a syntax error is raised, and it should follow the function name
+                elementToPrecedeType = elementToPrecedeType.getNextSibling();
+            }
+            Pair<PsiElement, PsiElement> typeWhiteSpaceAndColon = JetPsiFactory.createTypeWhiteSpaceAndColon(project, type.toString());
+            element.addRangeAfter(typeWhiteSpaceAndColon.first, typeWhiteSpaceAndColon.second, elementToPrecedeType);
+        }
+    }
+
+    @NotNull
+    public static JetIntentionActionFactory createFactoryForComponentFunctionReturnTypeMismatch() {
+        return new JetIntentionActionFactory() {
+            @Nullable
+            @Override
+            public IntentionAction createAction(Diagnostic diagnostic) {
+                String componentName = ((DiagnosticWithParameters3<JetExpression, Name, JetType, JetType>) diagnostic).getA().getName();
+                int componentIndex = Integer.valueOf(componentName.substring(DescriptorResolver.COMPONENT_FUNCTION_NAME_PREFIX.length()));
+                JetMultiDeclaration multiDeclaration = QuickFixUtil.getParentElementOfType(diagnostic, JetMultiDeclaration.class);
+                assert multiDeclaration != null : "COMPONENT_FUNCTION_RETURN_TYPE_MISMATCH reported on expression that is not within any multi declaration";
+                JetMultiDeclarationEntry entry = multiDeclaration.getEntries().get(componentIndex - 1);
+                BindingContext context = AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) multiDeclaration.getContainingFile().getContainingFile()).getBindingContext();
+                ResolvedCall<FunctionDescriptor> resolvedCall = context.get(BindingContext.COMPONENT_RESOLVED_CALL, entry);
+                if (resolvedCall == null) return null;
+                JetFunction componentFunction = (JetFunction) BindingContextUtils.descriptorToDeclaration(context, resolvedCall.getCandidateDescriptor());
+                JetType expectedType = context.get(BindingContext.TYPE, entry.getTypeRef());
+                if (componentFunction != null && expectedType != null) {
+                    return new ChangeFunctionReturnTypeFix(componentFunction, expectedType);
+                }
+                else return null;
+            }
+        };
+    }
+
+    @NotNull
+    public static JetIntentionActionFactory createFactoryForHasNextFunctionTypeMismatch() {
+        return new JetIntentionActionFactory() {
+            @Nullable
+            @Override
+            public IntentionAction createAction(Diagnostic diagnostic) {
+                JetExpression expression = QuickFixUtil.getParentElementOfType(diagnostic, JetExpression.class);
+                assert expression != null : "HAS_NEXT_FUNCTION_TYPE_MISMATCH reported on element that is not within any expression";
+                BindingContext context = AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) expression.getContainingFile()).getBindingContext();
+                ResolvedCall<FunctionDescriptor> resolvedCall = context.get(BindingContext.LOOP_RANGE_HAS_NEXT_RESOLVED_CALL, expression);
+                if (resolvedCall == null) return null;
+                JetFunction hasNextFunction = (JetFunction) BindingContextUtils.descriptorToDeclaration(context, resolvedCall.getCandidateDescriptor());
+                if (hasNextFunction != null) {
+                    return new ChangeFunctionReturnTypeFix(hasNextFunction, KotlinBuiltIns.getInstance().getBooleanType());
+                }
+                else return null;
+            }
+        };
+    }
+}

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

@@ -212,6 +212,8 @@ public class QuickFixes {
 
         factories.put(RETURN_TYPE_MISMATCH_ON_OVERRIDE, ChangeReturnTypeToMatchOverriddenMethodFix.createFactory());
         factories.put(PROPERTY_TYPE_MISMATCH_ON_OVERRIDE, ChangePropertyTypeToMatchOverriddenPropertyFix.createFactory());
+        factories.put(COMPONENT_FUNCTION_RETURN_TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForComponentFunctionReturnTypeMismatch());
+        factories.put(HAS_NEXT_FUNCTION_TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForHasNextFunctionTypeMismatch());
 
         factories.put(PLATFORM_CLASS_MAPPED_TO_KOTLIN, MapPlatformClassToKotlinFix.createFactory());
     }

idea/testData/quickfix/typeMismatch/afterChangeReturnTypeWhenValueParameterListIsAbsent.kt

@@ -0,0 +1,10 @@
+// "Change 'A.hasNext' function return type to 'Boolean'" "true"
+abstract class A {
+    abstract fun hasNext: Boolean
+    abstract fun next(): Int
+    abstract fun iterator(): A
+}
+
+fun test(notRange: A) {
+    for (i in notRange<caret>) {}
+}

idea/testData/quickfix/typeMismatch/afterComponentFunctionReturnTypeMismatch1.kt

@@ -0,0 +1,9 @@
+// "Change 'A.component1' function return type to 'Int'" "true"
+abstract class A {
+    abstract fun component1(): Int
+    abstract fun component2(): Int
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Int) = a<caret>
+}

idea/testData/quickfix/typeMismatch/afterComponentFunctionReturnTypeMismatch2.kt

@@ -0,0 +1,9 @@
+// "Change 'A.component2' function return type to 'Int'" "true"
+abstract class A {
+    abstract fun component1(): Int
+    abstract fun component2(): Int
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Int) = a<caret>
+}

idea/testData/quickfix/typeMismatch/afterComponentFunctionReturnTypeMismatch3.kt

@@ -0,0 +1,9 @@
+// "Remove explicitly specified return type in 'A.component2' function" "true"
+abstract class A {
+    abstract fun component1(): Int
+    abstract fun component2()
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Unit) = a<caret>
+}

idea/testData/quickfix/typeMismatch/afterComponentFunctionReturnTypeMismatch4.kt

@@ -0,0 +1,10 @@
+// "Change 'A.component2' function return type to 'Unit'" "true"
+// ERROR: <html>Type mismatch.<table><tr><td>Required:</td><td>jet.Unit</td></tr><tr><td>Found:</td><td>jet.Int</td></tr></table></html>
+abstract class A {
+    abstract fun component1(): Int
+    fun component2(): Unit = 42
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Unit) = a<caret>
+}

idea/testData/quickfix/typeMismatch/afterHasNextFunctionReturnTypeMismatch.kt

@@ -0,0 +1,10 @@
+// "Change 'A.hasNext' function return type to 'Boolean'" "true"
+abstract class A {
+    abstract fun hasNext(): Boolean
+    abstract fun next(): Int
+    abstract fun iterator(): A
+}
+
+fun test(notRange: A) {
+    for (i in notRange<caret>) {}
+}

idea/testData/quickfix/typeMismatch/beforeChangeReturnTypeWhenValueParameterListIsAbsent.kt

@@ -0,0 +1,10 @@
+// "Change 'A.hasNext' function return type to 'Boolean'" "true"
+abstract class A {
+    abstract fun hasNext
+    abstract fun next(): Int
+    abstract fun iterator(): A
+}
+
+fun test(notRange: A) {
+    for (i in notRange<caret>) {}
+}

idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch1.kt

@@ -0,0 +1,9 @@
+// "Change 'A.component1' function return type to 'Int'" "true"
+abstract class A {
+    abstract fun component1()
+    abstract fun component2(): Int
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Int) = a<caret>
+}

idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch2.kt

@@ -0,0 +1,9 @@
+// "Change 'A.component2' function return type to 'Int'" "true"
+abstract class A {
+    abstract fun component1(): Int
+    abstract fun component2(): String
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Int) = a<caret>
+}

idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch3.kt

@@ -0,0 +1,9 @@
+// "Remove explicitly specified return type in 'A.component2' function" "true"
+abstract class A {
+    abstract fun component1(): Int
+    abstract fun component2(): Int
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Unit) = a<caret>
+}

idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch4.kt

@@ -0,0 +1,10 @@
+// "Change 'A.component2' function return type to 'Unit'" "true"
+// ERROR: <html>Type mismatch.<table><tr><td>Required:</td><td>jet.Unit</td></tr><tr><td>Found:</td><td>jet.Int</td></tr></table></html>
+abstract class A {
+    abstract fun component1(): Int
+    fun component2() = 42
+}
+
+fun foo(a: A) {
+    val (w: Int, x: Unit) = a<caret>
+}

idea/testData/quickfix/typeMismatch/beforeHasNextFunctionReturnTypeMismatch.kt

@@ -0,0 +1,10 @@
+// "Change 'A.hasNext' function return type to 'Boolean'" "true"
+abstract class A {
+    abstract fun hasNext(): Int
+    abstract fun next(): Int
+    abstract fun iterator(): A
+}
+
+fun test(notRange: A) {
+    for (i in notRange<caret>) {}
+}

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

@@ -31,7 +31,7 @@ import org.jetbrains.jet.plugin.quickfix.AbstractQuickFixTest;
 /** This class is generated by {@link org.jetbrains.jet.generators.tests.GenerateTests}. DO NOT MODIFY MANUALLY */
 @SuppressWarnings("all")
 @TestMetadata("idea/testData/quickfix")
-@InnerTestClasses({QuickFixTestGenerated.Abstract.class, QuickFixTestGenerated.AddStarProjections.class, QuickFixTestGenerated.AutoImports.class, QuickFixTestGenerated.CheckArguments.class, QuickFixTestGenerated.Expressions.class, QuickFixTestGenerated.Migration.class, QuickFixTestGenerated.Modifiers.class, QuickFixTestGenerated.Nullables.class, QuickFixTestGenerated.Override.class, QuickFixTestGenerated.PlatformClasses.class, QuickFixTestGenerated.Supercalls.class, QuickFixTestGenerated.SupertypeInitialization.class, QuickFixTestGenerated.TypeAddition.class, QuickFixTestGenerated.TypeImports.class, QuickFixTestGenerated.TypeProjection.class, QuickFixTestGenerated.UselessImports.class, QuickFixTestGenerated.Variables.class, QuickFixTestGenerated.When.class})
+@InnerTestClasses({QuickFixTestGenerated.Abstract.class, QuickFixTestGenerated.AddStarProjections.class, QuickFixTestGenerated.AutoImports.class, QuickFixTestGenerated.CheckArguments.class, QuickFixTestGenerated.Expressions.class, QuickFixTestGenerated.Migration.class, QuickFixTestGenerated.Modifiers.class, QuickFixTestGenerated.Nullables.class, QuickFixTestGenerated.Override.class, QuickFixTestGenerated.PlatformClasses.class, QuickFixTestGenerated.Supercalls.class, QuickFixTestGenerated.SupertypeInitialization.class, QuickFixTestGenerated.TypeAddition.class, QuickFixTestGenerated.TypeImports.class, QuickFixTestGenerated.TypeMismatch.class, QuickFixTestGenerated.TypeProjection.class, QuickFixTestGenerated.UselessImports.class, QuickFixTestGenerated.Variables.class, QuickFixTestGenerated.When.class})
 public class QuickFixTestGenerated extends AbstractQuickFixTest {
     public void testAllFilesPresentInQuickfix() throws Exception {
         JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/quickfix"), Pattern.compile("^before(\\w+)\\.kt$"), true);
@@ -967,6 +967,44 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
         
     }
     
+    @TestMetadata("idea/testData/quickfix/typeMismatch")
+    public static class TypeMismatch extends AbstractQuickFixTest {
+        public void testAllFilesPresentInTypeMismatch() throws Exception {
+            JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/quickfix/typeMismatch"), Pattern.compile("^before(\\w+)\\.kt$"), true);
+        }
+        
+        @TestMetadata("beforeChangeReturnTypeWhenValueParameterListIsAbsent.kt")
+        public void testChangeReturnTypeWhenValueParameterListIsAbsent() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeChangeReturnTypeWhenValueParameterListIsAbsent.kt");
+        }
+        
+        @TestMetadata("beforeComponentFunctionReturnTypeMismatch1.kt")
+        public void testComponentFunctionReturnTypeMismatch1() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch1.kt");
+        }
+        
+        @TestMetadata("beforeComponentFunctionReturnTypeMismatch2.kt")
+        public void testComponentFunctionReturnTypeMismatch2() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch2.kt");
+        }
+        
+        @TestMetadata("beforeComponentFunctionReturnTypeMismatch3.kt")
+        public void testComponentFunctionReturnTypeMismatch3() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch3.kt");
+        }
+        
+        @TestMetadata("beforeComponentFunctionReturnTypeMismatch4.kt")
+        public void testComponentFunctionReturnTypeMismatch4() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeComponentFunctionReturnTypeMismatch4.kt");
+        }
+        
+        @TestMetadata("beforeHasNextFunctionReturnTypeMismatch.kt")
+        public void testHasNextFunctionReturnTypeMismatch() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeHasNextFunctionReturnTypeMismatch.kt");
+        }
+        
+    }
+    
     @TestMetadata("idea/testData/quickfix/typeProjection")
     public static class TypeProjection extends AbstractQuickFixTest {
         public void testAllFilesPresentInTypeProjection() throws Exception {
@@ -1198,6 +1236,7 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
         suite.addTestSuite(SupertypeInitialization.class);
         suite.addTestSuite(TypeAddition.class);
         suite.addTestSuite(TypeImports.class);
+        suite.addTestSuite(TypeMismatch.class);
         suite.addTestSuite(TypeProjection.class);
         suite.addTestSuite(UselessImports.class);
         suite.addTest(Variables.innerSuite());