Change return type quick-fix: Add support of constant expressions

#KT-6260 Fixed
 #KT-6404 Fixed

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

@@ -17,37 +17,76 @@
 package org.jetbrains.jet.plugin.quickfix;
 
 import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.CallableDescriptor;
 import org.jetbrains.jet.lang.diagnostics.Diagnostic;
+import org.jetbrains.jet.lang.diagnostics.DiagnosticWithParameters1;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticWithParameters2;
 import org.jetbrains.jet.lang.diagnostics.Errors;
+import org.jetbrains.jet.lang.diagnostics.rendering.DefaultErrorMessages;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.calls.callUtil.CallUtilPackage;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.plugin.caches.resolve.ResolvePackage;
+import org.jetbrains.jet.plugin.util.UtilPackage;
 
+import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 
 //TODO: should use change signature to deal with cases of multiple overridden descriptors
 public class QuickFixFactoryForTypeMismatchError extends JetIntentionActionsFactory {
+    private final static Logger LOG = Logger.getInstance(QuickFixFactoryForTypeMismatchError.class);
+
     @NotNull
     @Override
     protected List<IntentionAction> doCreateActions(@NotNull Diagnostic diagnostic) {
         List<IntentionAction> actions = new LinkedList<IntentionAction>();
 
-        DiagnosticWithParameters2<JetExpression, JetType, JetType> diagnosticWithParameters = Errors.TYPE_MISMATCH.cast(diagnostic);
-        JetExpression expression = diagnosticWithParameters.getPsiElement();
-        JetType expectedType = diagnosticWithParameters.getA();
-        JetType expressionType = diagnosticWithParameters.getB();
         BindingContext context = ResolvePackage.analyzeFully((JetFile) diagnostic.getPsiFile());
 
+        PsiElement diagnosticElement = diagnostic.getPsiElement();
+        if (!(diagnosticElement instanceof JetExpression)) {
+            LOG.error("Unexpected element: " + diagnosticElement.getText());
+            return Collections.emptyList();
+        }
+
+        JetExpression expression = (JetExpression) diagnosticElement;
+
+        JetType expectedType;
+        JetType expressionType;
+        if (diagnostic.getFactory() == Errors.TYPE_MISMATCH) {
+            DiagnosticWithParameters2<JetExpression, JetType, JetType> diagnosticWithParameters = Errors.TYPE_MISMATCH.cast(diagnostic);
+            expectedType = diagnosticWithParameters.getA();
+            expressionType = diagnosticWithParameters.getB();
+        }
+        else if (diagnostic.getFactory() == Errors.NULL_FOR_NONNULL_TYPE) {
+            DiagnosticWithParameters1<JetConstantExpression, JetType> diagnosticWithParameters =
+                    Errors.NULL_FOR_NONNULL_TYPE.cast(diagnostic);
+            expectedType = diagnosticWithParameters.getA();
+            expressionType = UtilPackage.makeNullable(expectedType);
+        }
+        else if (diagnostic.getFactory() == Errors.CONSTANT_EXPECTED_TYPE_MISMATCH) {
+            DiagnosticWithParameters2<JetConstantExpression, String, JetType> diagnosticWithParameters =
+                    Errors.CONSTANT_EXPECTED_TYPE_MISMATCH.cast(diagnostic);
+            expectedType = diagnosticWithParameters.getB();
+            expressionType = context.get(BindingContext.EXPRESSION_TYPE, expression);
+            if (expressionType == null) {
+                LOG.error("No type inferred: " + expression.getText());
+                return Collections.emptyList();
+            }
+        }
+        else {
+            LOG.error("Unexpected diagnostic: " + DefaultErrorMessages.RENDERER.render(diagnostic));
+            return Collections.emptyList();
+        }
+
         // We don't want to cast a cast or type-asserted expression:
         if (!(expression instanceof JetBinaryExpressionWithTypeRHS) && !(expression.getParent() instanceof  JetBinaryExpressionWithTypeRHS)) {
             actions.add(new CastExpressionFix(expression, expectedType));

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

@@ -254,7 +254,10 @@ public class QuickFixRegistrar {
 
         QuickFixes.factories.put(FUNCTION_EXPECTED, CreateInvokeFunctionActionFactory.INSTANCE$);
 
-        QuickFixes.factories.put(TYPE_MISMATCH, new QuickFixFactoryForTypeMismatchError());
+        QuickFixFactoryForTypeMismatchError factoryForTypeMismatchError = new QuickFixFactoryForTypeMismatchError();
+        QuickFixes.factories.put(TYPE_MISMATCH, factoryForTypeMismatchError);
+        QuickFixes.factories.put(NULL_FOR_NONNULL_TYPE, factoryForTypeMismatchError);
+        QuickFixes.factories.put(CONSTANT_EXPECTED_TYPE_MISMATCH, factoryForTypeMismatchError);
 
         QuickFixes.factories.put(SMARTCAST_IMPOSSIBLE, CastExpressionFix.createFactoryForSmartCastImpossible());
 

idea/testData/quickfix/typeMismatch/afterChangeReturnTypeToSpecificNullable.kt

@@ -0,0 +1,6 @@
+// "Change 'foo' function return type to 'Int?'" "true"
+
+fun foo(): Int? {
+    val n: Int? = 1
+    return n
+}

idea/testData/quickfix/typeMismatch/afterConstantTypeMismatch.kt

@@ -0,0 +1,5 @@
+// "Change 'foo' function return type to 'Int'" "true"
+
+fun foo(): Int {
+    return <caret>1
+}

idea/testData/quickfix/typeMismatch/afterMakeReturnTypeNullable.kt

@@ -0,0 +1,5 @@
+// "Change 'foo' function return type to 'String?'" "true"
+
+fun foo(): String? {
+    return null
+}

idea/testData/quickfix/typeMismatch/beforeChangeReturnTypeToSpecificNullable.kt

@@ -0,0 +1,6 @@
+// "Change 'foo' function return type to 'Int?'" "true"
+
+fun foo(): String {
+    val n: Int? = 1
+    return <caret>n
+}

idea/testData/quickfix/typeMismatch/beforeConstantTypeMismatch.kt

@@ -0,0 +1,5 @@
+// "Change 'foo' function return type to 'Int'" "true"
+
+fun foo(): String {
+    return <caret>1
+}

idea/testData/quickfix/typeMismatch/beforeMakeReturnTypeNullable.kt

@@ -0,0 +1,5 @@
+// "Change 'foo' function return type to 'String?'" "true"
+
+fun foo(): String {
+    return <caret>null
+}

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

@@ -4166,6 +4166,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("beforeChangeReturnTypeToSpecificNullable.kt")
+        public void testChangeReturnTypeToSpecificNullable() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeChangeReturnTypeToSpecificNullable.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("beforeChangeReturnTypeWhenFunctionNameIsMissing.kt")
         public void testChangeReturnTypeWhenFunctionNameIsMissing() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeChangeReturnTypeWhenFunctionNameIsMissing.kt");
@@ -4184,6 +4190,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("beforeConstantTypeMismatch.kt")
+        public void testConstantTypeMismatch() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeConstantTypeMismatch.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("beforeDontChangeOverriddenPropertyTypeToErrorType.kt")
         public void testDontChangeOverriddenPropertyTypeToErrorType() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeDontChangeOverriddenPropertyTypeToErrorType.kt");
@@ -4214,6 +4226,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("beforeMakeReturnTypeNullable.kt")
+        public void testMakeReturnTypeNullable() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeMakeReturnTypeNullable.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("beforeNoReturnInFunctionWithBlockBody.kt")
         public void testNoReturnInFunctionWithBlockBody() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/beforeNoReturnInFunctionWithBlockBody.kt");