ChangeFunctionLiteralReturnTypeFix for EXPECTED/ASSIGNMENT_TYPE_MISMATCH errors

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

@@ -59,6 +59,7 @@ 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}''
 change.no.name.function.return.type=Change function return type to ''{0}''
+change.function.literal.return.type=Change function literal return type to ''{0}''
 remove.function.return.type=Remove explicitly specified return type in ''{0}'' function
 remove.no.name.function.return.type=Remove explicitly specified function return type
 change.element.type=Change ''{0}'' type to ''{1}''

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

@@ -0,0 +1,140 @@
+/*
+ * 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.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.descriptors.ClassDescriptor;
+import org.jetbrains.jet.lang.diagnostics.Diagnostic;
+import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.types.JetType;
+import org.jetbrains.jet.lang.types.TypeProjection;
+import org.jetbrains.jet.lang.types.TypeUtils;
+import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
+import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
+import org.jetbrains.jet.plugin.JetBundle;
+import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache;
+import org.jetbrains.jet.renderer.DescriptorRenderer;
+
+import java.util.LinkedList;
+import java.util.List;
+
+public class ChangeFunctionLiteralReturnTypeFix extends JetIntentionAction<JetFunctionLiteralExpression> {
+    private final String renderedType;
+    private final JetTypeReference functionLiteralReturnTypeRef;
+    private IntentionAction appropriateQuickFix = null;
+
+    public ChangeFunctionLiteralReturnTypeFix(@NotNull JetFunctionLiteralExpression element, @NotNull JetType type) {
+        super(element);
+        renderedType = DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(type);
+        functionLiteralReturnTypeRef = element.getFunctionLiteral().getReturnTypeRef();
+
+        BindingContext context = AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) element.getContainingFile()).getBindingContext();
+        JetType functionLiteralType = context.get(BindingContext.EXPRESSION_TYPE, element);
+        assert functionLiteralType != null : "Type of function literal not available in binding context";
+
+        ClassDescriptor functionClass = KotlinBuiltIns.getInstance().getFunction(functionLiteralType.getArguments().size() - 1);
+        List<JetType> functionClassTypeParameters = new LinkedList<JetType>();
+        for (TypeProjection typeProjection: functionLiteralType.getArguments()) {
+            functionClassTypeParameters.add(typeProjection.getType());
+        }
+        // Replacing return type:
+        functionClassTypeParameters.remove(functionClassTypeParameters.size() - 1);
+        functionClassTypeParameters.add(type);
+        JetType eventualFunctionLiteralType = TypeUtils.substituteParameters(functionClass, functionClassTypeParameters);
+
+        JetProperty correspondingProperty = PsiTreeUtil.getParentOfType(element, JetProperty.class);
+        if (correspondingProperty != null && QuickFixUtil.canEvaluateTo(correspondingProperty.getInitializer(), element)) {
+            JetTypeReference correspondingPropertyTypeRef = correspondingProperty.getTypeRef();
+            JetType propertyType = context.get(BindingContext.TYPE, correspondingPropertyTypeRef);
+            if (propertyType != null && !JetTypeChecker.INSTANCE.isSubtypeOf(eventualFunctionLiteralType, propertyType)) {
+                appropriateQuickFix = new ChangeVariableTypeFix(correspondingProperty, eventualFunctionLiteralType);
+            }
+            return;
+        }
+
+        JetParameter correspondingParameter = QuickFixUtil.getFunctionParameterCorrespondingToFunctionLiteralPassedOutsideArgumentList(element);
+        if (correspondingParameter != null) {
+            JetTypeReference correspondingParameterTypeRef = correspondingParameter.getTypeReference();
+            JetType parameterType = context.get(BindingContext.TYPE, correspondingParameterTypeRef);
+            if (parameterType != null && !JetTypeChecker.INSTANCE.isSubtypeOf(eventualFunctionLiteralType, parameterType)) {
+                appropriateQuickFix = new ChangeFunctionParameterTypeFix(correspondingParameter, eventualFunctionLiteralType);
+            }
+            return;
+        }
+
+        JetFunction parentFunction = PsiTreeUtil.getParentOfType(element, JetFunction.class, true);
+        if (parentFunction != null && QuickFixUtil.canFunctionReturnExpression(parentFunction, element)) {
+            JetTypeReference parentFunctionReturnTypeRef = parentFunction.getReturnTypeRef();
+            JetType parentFunctionReturnType = context.get(BindingContext.TYPE, parentFunctionReturnTypeRef);
+            if (parentFunctionReturnType != null && !JetTypeChecker.INSTANCE.isSubtypeOf(eventualFunctionLiteralType, parentFunctionReturnType)) {
+                appropriateQuickFix = new ChangeFunctionReturnTypeFix(parentFunction, eventualFunctionLiteralType);
+            }
+        }
+    }
+
+    @NotNull
+    @Override
+    public String getText() {
+        if (appropriateQuickFix != null) {
+            return appropriateQuickFix.getText();
+        }
+        return JetBundle.message("change.function.literal.return.type", renderedType);
+    }
+
+    @NotNull
+    @Override
+    public String getFamilyName() {
+        return JetBundle.message("change.type.family");
+    }
+
+    @Override
+    public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+        return super.isAvailable(project, editor, file) &&
+            (functionLiteralReturnTypeRef != null || (appropriateQuickFix != null && appropriateQuickFix.isAvailable(project, editor, file)));
+    }
+
+    @Override
+    public void invoke(@NotNull Project project, Editor editor, JetFile file) throws IncorrectOperationException {
+        if (functionLiteralReturnTypeRef != null) {
+            functionLiteralReturnTypeRef.replace(JetPsiFactory.createType(project, renderedType));
+        }
+        if (appropriateQuickFix != null && appropriateQuickFix.isAvailable(project, editor, file)) {
+            appropriateQuickFix.invoke(project, editor, file);
+        }
+    }
+
+    @NotNull
+    public static JetIntentionActionFactory createFactoryForExpectedOrAssignmentTypeMismatch() {
+        return new JetIntentionActionFactory() {
+            @Nullable
+            @Override
+            public IntentionAction createAction(Diagnostic diagnostic) {
+                JetFunctionLiteralExpression functionLiteralExpression = QuickFixUtil.getParentElementOfType(diagnostic, JetFunctionLiteralExpression.class);
+                assert functionLiteralExpression != null : "ASSIGNMENT/EXPECTED_TYPE_MISMATCH reported outside any function literal";
+                return new ChangeFunctionLiteralReturnTypeFix(functionLiteralExpression, KotlinBuiltIns.getInstance().getUnitType());
+            }
+        };
+    }
+}

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

@@ -149,4 +149,29 @@ public class QuickFixUtil {
             return parameterList.getParameters().get(position);
         }
     }
+
+    private static boolean equalOrLastInThenOrElse(JetExpression thenOrElse, JetExpression expression) {
+        if (thenOrElse == expression) return true;
+        return thenOrElse instanceof JetBlockExpression && expression.getParent() == thenOrElse &&
+               PsiTreeUtil.getNextSiblingOfType(expression, JetExpression.class) == null;
+    }
+
+    public static boolean canEvaluateTo(JetExpression parent, JetExpression child) {
+        if (parent == null || child == null) {
+            return false;
+        }
+        while (parent != child) {
+            if (child.getParent() instanceof JetParenthesizedExpression) {
+                child = (JetExpression) child.getParent();
+                continue;
+            }
+            JetIfExpression jetIfExpression = PsiTreeUtil.getParentOfType(child, JetIfExpression.class, true);
+            if (jetIfExpression == null) return false;
+            if (!equalOrLastInThenOrElse(jetIfExpression.getThen(), child) && !equalOrLastInThenOrElse(jetIfExpression.getElse(), child)) {
+                return false;
+            }
+            child = jetIfExpression;
+        }
+        return true;
+    }
 }

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

@@ -236,6 +236,10 @@ public class QuickFixes {
         factories.put(EXPECTED_PARAMETER_TYPE_MISMATCH, ChangeTypeFix.createFactoryForExpectedParameterTypeMismatch());
         factories.put(EXPECTED_RETURN_TYPE_MISMATCH, ChangeTypeFix.createFactoryForExpectedReturnTypeMismatch());
 
+        JetIntentionActionFactory changeFunctionLiteralReturnTypeFix = ChangeFunctionLiteralReturnTypeFix.createFactoryForExpectedOrAssignmentTypeMismatch();
+        factories.put(EXPECTED_TYPE_MISMATCH, changeFunctionLiteralReturnTypeFix);
+        factories.put(ASSIGNMENT_TYPE_MISMATCH, changeFunctionLiteralReturnTypeFix);
+
         factories.put(TYPE_MISMATCH, ChangeFunctionParameterTypeFix.createFactory());
 
         factories.put(AUTOCAST_IMPOSSIBLE, CastExpressionFix.createFactoryForAutoCastImpossible());

idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/afterAssignmentTypeMismatch.kt

@@ -0,0 +1,7 @@
+// "Change 'f' type to '() -> Unit'" "true"
+fun foo() {
+    val f: () -> Unit = {
+        var x = 1
+        x += 21<caret>
+    }
+}

idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/afterExpectedTypeMismatch.kt

@@ -0,0 +1,7 @@
+// "Change 'bar' type to '(() -> Long) -> Unit'" "true"
+fun foo() {
+    val bar: (() -> Long) -> Unit = {
+        (f: () -> Long): Unit ->
+        var x = 5<caret>
+    }
+}

idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeAssignmentTypeMismatch.kt

@@ -0,0 +1,7 @@
+// "Change 'f' type to '() -> Unit'" "true"
+fun foo() {
+    val f: () -> Int = {
+        var x = 1
+        x += 21<caret>
+    }
+}

idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeExpectedTypeMismatch.kt

@@ -0,0 +1,7 @@
+// "Change 'bar' type to '(() -> Long) -> Unit'" "true"
+fun foo() {
+    val bar: () -> Double = {
+        (f: () -> Long): String ->
+        var x = 5<caret>
+    }
+}

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

@@ -1416,11 +1416,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                 JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression"), Pattern.compile("^before(\\w+)\\.kt$"), true);
             }
             
+            @TestMetadata("beforeAssignmentTypeMismatch.kt")
+            public void testAssignmentTypeMismatch() throws Exception {
+                doTest("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeAssignmentTypeMismatch.kt");
+            }
+            
             @TestMetadata("beforeChangeFunctionReturnTypeToFunctionType.kt")
             public void testChangeFunctionReturnTypeToFunctionType() throws Exception {
                 doTest("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeChangeFunctionReturnTypeToFunctionType.kt");
             }
             
+            @TestMetadata("beforeExpectedTypeMismatch.kt")
+            public void testExpectedTypeMismatch() throws Exception {
+                doTest("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeExpectedTypeMismatch.kt");
+            }
+            
             @TestMetadata("beforeTypeMismatchInInitializer.kt")
             public void testTypeMismatchInInitializer() throws Exception {
                 doTest("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeTypeMismatchInInitializer.kt");