From 9b11db7a742665f31331eb01c96e5c84c14dce1c Mon Sep 17 00:00:00 2001
From: Wojciech Lopata <wlopata@gmail.com>
Date: Wed, 24 Apr 2013 14:32:44 +0200
Subject: [PATCH] ChangeFunctionReturnTypeFix for TYPE_MISMATCH error

---
 .../quickfix/ChangeFunctionReturnTypeFix.java | 20 +++++++++++++++++++
 .../jet/plugin/quickfix/QuickFixes.java       |  1 +
 .../afterTypeMismatchInInitializer.kt         |  2 ++
 .../afterTypeMismatchInReturnStatement.kt     |  4 ++++
 .../beforeTypeMismatchInInitializer.kt        |  2 ++
 .../beforeTypeMismatchInReturnStatement.kt    |  4 ++++
 .../quickfix/QuickFixTestGenerated.java       | 10 ++++++++++
 7 files changed, 43 insertions(+)
 create mode 100644 idea/testData/quickfix/typeMismatch/afterTypeMismatchInInitializer.kt
 create mode 100644 idea/testData/quickfix/typeMismatch/afterTypeMismatchInReturnStatement.kt
 create mode 100644 idea/testData/quickfix/typeMismatch/beforeTypeMismatchInInitializer.kt
 create mode 100644 idea/testData/quickfix/typeMismatch/beforeTypeMismatchInReturnStatement.kt

diff --git a/idea/src/org/jetbrains/jet/plugin/quickfix/ChangeFunctionReturnTypeFix.java b/idea/src/org/jetbrains/jet/plugin/quickfix/ChangeFunctionReturnTypeFix.java
index ea34a1dbe19..246d3c322e5 100644
--- a/idea/src/org/jetbrains/jet/plugin/quickfix/ChangeFunctionReturnTypeFix.java
+++ b/idea/src/org/jetbrains/jet/plugin/quickfix/ChangeFunctionReturnTypeFix.java
@@ -185,4 +185,24 @@ public class ChangeFunctionReturnTypeFix extends JetIntentionAction<JetNamedFunc
             }
         };
     }
+
+    @NotNull
+    public static JetIntentionActionFactory createFactoryForTypeMismatch() {
+        return new JetIntentionActionFactory() {
+            @Nullable
+            @Override
+            public IntentionAction createAction(Diagnostic diagnostic) {
+                JetExpression expression = (JetExpression) diagnostic.getPsiElement();
+                JetNamedFunction function = QuickFixUtil.getParentElementOfType(diagnostic, JetNamedFunction.class);
+
+                if (function != null && (function.getInitializer() == expression || expression.getParent() instanceof JetReturnExpression)) {
+                    BindingContext context = AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) diagnostic.getPsiFile()).getBindingContext();
+                    JetType type = context.get(BindingContext.EXPRESSION_TYPE, expression);
+                    assert type != null : "Expression type mismatch, but expression has no type";
+                    return new ChangeFunctionReturnTypeFix(function, type);
+                }
+                return null;
+            }
+        };
+    }
 }
diff --git a/idea/src/org/jetbrains/jet/plugin/quickfix/QuickFixes.java b/idea/src/org/jetbrains/jet/plugin/quickfix/QuickFixes.java
index a4905cc5256..d67c7059070 100644
--- a/idea/src/org/jetbrains/jet/plugin/quickfix/QuickFixes.java
+++ b/idea/src/org/jetbrains/jet/plugin/quickfix/QuickFixes.java
@@ -222,6 +222,7 @@ public class QuickFixes {
         factories.put(COMPONENT_FUNCTION_RETURN_TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForComponentFunctionReturnTypeMismatch());
         factories.put(HAS_NEXT_FUNCTION_TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForHasNextFunctionTypeMismatch());
         factories.put(COMPARE_TO_TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForCompareToTypeMismatch());
+        factories.put(TYPE_MISMATCH, ChangeFunctionReturnTypeFix.createFactoryForTypeMismatch());
 
         factories.put(EXPECTED_PARAMETER_TYPE_MISMATCH, ChangeTypeFix.createFactoryForExpectedParameterTypeMismatch());
         factories.put(EXPECTED_RETURN_TYPE_MISMATCH, ChangeTypeFix.createFactoryForExpectedReturnTypeMismatch());
diff --git a/idea/testData/quickfix/typeMismatch/afterTypeMismatchInInitializer.kt b/idea/testData/quickfix/typeMismatch/afterTypeMismatchInInitializer.kt
new file mode 100644
index 00000000000..4528d211219
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/afterTypeMismatchInInitializer.kt
@@ -0,0 +1,2 @@
+// "Change 'foo' function return type to 'String'" "true"
+fun foo(): String = <caret>""
\ No newline at end of file
diff --git a/idea/testData/quickfix/typeMismatch/afterTypeMismatchInReturnStatement.kt b/idea/testData/quickfix/typeMismatch/afterTypeMismatchInReturnStatement.kt
new file mode 100644
index 00000000000..eb37598a97b
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/afterTypeMismatchInReturnStatement.kt
@@ -0,0 +1,4 @@
+// "Change 'foo' function return type to 'String'" "true"
+fun foo(): String {
+    return ""<caret>
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInInitializer.kt b/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInInitializer.kt
new file mode 100644
index 00000000000..7fc2c617bd9
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInInitializer.kt
@@ -0,0 +1,2 @@
+// "Change 'foo' function return type to 'String'" "true"
+fun foo(): Int = <caret>""
\ No newline at end of file
diff --git a/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInReturnStatement.kt b/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInReturnStatement.kt
new file mode 100644
index 00000000000..71a0a7aa945
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/beforeTypeMismatchInReturnStatement.kt
@@ -0,0 +1,4 @@
+// "Change 'foo' function return type to 'String'" "true"
+fun foo() {
+    return ""<caret>
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/jet/plugin/quickfix/QuickFixTestGenerated.java b/idea/tests/org/jetbrains/jet/plugin/quickfix/QuickFixTestGenerated.java
index 05a974141e4..e26c3ac86b9 100644
--- a/idea/tests/org/jetbrains/jet/plugin/quickfix/QuickFixTestGenerated.java
+++ b/idea/tests/org/jetbrains/jet/plugin/quickfix/QuickFixTestGenerated.java
@@ -1063,6 +1063,16 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest("idea/testData/quickfix/typeMismatch/beforeReturnTypeMismatch.kt");
         }
         
+        @TestMetadata("beforeTypeMismatchInInitializer.kt")
+        public void testTypeMismatchInInitializer() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeTypeMismatchInInitializer.kt");
+        }
+        
+        @TestMetadata("beforeTypeMismatchInReturnStatement.kt")
+        public void testTypeMismatchInReturnStatement() throws Exception {
+            doTest("idea/testData/quickfix/typeMismatch/beforeTypeMismatchInReturnStatement.kt");
+        }
+        
     }
     
     @TestMetadata("idea/testData/quickfix/typeProjection")