Quick-Fixes: Support "Change function return type" quick-fix for non-local returns

#KT-6910 Fixed

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContextUtils.kt

@@ -26,7 +26,9 @@ import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.JetDeclarationWithBody
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.psi.JetCallableDeclaration
 import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 
@@ -43,6 +45,10 @@ public fun JetReturnExpression.getTargetFunctionDescriptor(context: BindingConte
             .firstOrNull()
 }
 
+public fun JetReturnExpression.getTargetFunction(context: BindingContext): JetCallableDeclaration? {
+    return getTargetFunctionDescriptor(context)?.let { DescriptorToSourceUtils.descriptorToDeclaration(it) as? JetCallableDeclaration }
+}
+
 public fun JetExpression.isUsedAsExpression(context: BindingContext): Boolean = context[BindingContext.USED_AS_EXPRESSION, this]!!
 public fun JetExpression.isUsedAsStatement(context: BindingContext): Boolean = !isUsedAsExpression(context)
 

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixFactoryForTypeMismatchError.java

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.idea.util.UtilPackage;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.bindingContextUtil.BindingContextUtilPackage;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilPackage;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.types.JetType;
@@ -102,10 +103,15 @@ public class QuickFixFactoryForTypeMismatchError extends JetIntentionActionsFact
             }
         }
 
+        PsiElement expressionParent = expression.getParent();
+
         // Mismatch in returned expression:
-        JetFunction function = PsiTreeUtil.getParentOfType(expression, JetFunction.class, true);
-        if (function != null && QuickFixUtil.canFunctionOrGetterReturnExpression(function, expression)) {
-            actions.add(new ChangeFunctionReturnTypeFix(function, expressionType));
+
+        JetCallableDeclaration function = expressionParent instanceof JetReturnExpression
+                               ? BindingContextUtilPackage.getTargetFunction((JetReturnExpression) expressionParent, context)
+                               : PsiTreeUtil.getParentOfType(expression, JetFunction.class, true);
+        if (function instanceof JetFunction && QuickFixUtil.canFunctionOrGetterReturnExpression(function, expression)) {
+            actions.add(new ChangeFunctionReturnTypeFix((JetFunction) function, expressionType));
         }
 
         // Fixing overloaded operators:

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

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

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

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

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

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

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

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

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -4954,6 +4954,18 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("beforeNonLocalReturn.kt")
+            public void testNonLocalReturn() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeNonLocalReturn.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("beforeNonLocalReturnWithLabel.kt")
+            public void testNonLocalReturnWithLabel() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforeNonLocalReturnWithLabel.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("beforePropertyGetterInitializerTypeMismatch.kt")
             public void testPropertyGetterInitializerTypeMismatch() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/typeMismatch/typeMismatchOnReturnedExpression/beforePropertyGetterInitializerTypeMismatch.kt");