From 3b3bf18b699317f4218aee19fb7b4462e3dcdc8b Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Wed, 19 Aug 2015 12:27:28 +0200
Subject: [PATCH] suggest to fix return type based on args of method call in
 return stmt (IDEA-140894)

---
 .../impl/analysis/HighlightMethodUtil.java    | 33 +++++++++++++++++++
 .../methodReturn/afterInferFromArgs.java      | 12 +++++++
 .../methodReturn/beforeInferFromArgs.java     | 10 ++++++
 3 files changed, 55 insertions(+)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/afterInferFromArgs.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/beforeInferFromArgs.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightMethodUtil.java b/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightMethodUtil.java
index c8522b286307..cfd36c80aef7 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightMethodUtil.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightMethodUtil.java
@@ -41,6 +41,8 @@ import com.intellij.psi.infos.MethodCandidateInfo;
 import com.intellij.psi.util.*;
 import com.intellij.refactoring.util.RefactoringChangeUtil;
 import com.intellij.ui.ColorUtil;
+import com.intellij.util.Function;
+import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MostlySingularMultiMap;
 import com.intellij.util.ui.UIUtil;
 import com.intellij.xml.util.XmlStringUtil;
@@ -400,6 +402,7 @@ public class HighlightMethodUtil {
             .description(description).escapedToolTip(toolTip).navigationShift(navigationShift).create();
           if (highlightInfo != null) {
             registerMethodCallIntentions(highlightInfo, methodCall, list, resolveHelper);
+            registerMethodReturnFixAction(highlightInfo, candidateInfo, methodCall, resolveHelper);
           }
         }
         else {
@@ -435,6 +438,36 @@ public class HighlightMethodUtil {
     return highlightInfo;
   }
 
+  private static void registerMethodReturnFixAction(HighlightInfo highlightInfo,
+                                                    MethodCandidateInfo candidate,
+                                                    PsiMethodCallExpression methodCall,
+                                                    PsiResolveHelper resolveHelper) {
+    if (methodCall.getParent() instanceof PsiReturnStatement) {
+      final PsiMethod containerMethod = PsiTreeUtil.getParentOfType(methodCall, PsiMethod.class, true, PsiLambdaExpression.class);
+      if (containerMethod != null) {
+        final PsiMethod method = candidate.getElement();
+        final List<PsiType> list = ContainerUtil.map(method.getParameterList().getParameters(),
+                                                     new Function<PsiParameter, PsiType>() {
+                                                       @Override
+                                                       public PsiType fun(PsiParameter parameter) {
+                                                         return parameter.getType();
+                                                       }
+                                                     });
+        PsiType[] leftTypes = list.toArray(new PsiType[list.size()]);
+        final PsiSubstitutor substitutor = resolveHelper
+          .inferTypeArguments(method.getTypeParameters(), leftTypes, methodCall.getArgumentList().getExpressionTypes(),
+                              PsiUtil.getLanguageLevel(methodCall));
+        PsiType methodCallTypeByArgs = substitutor.substitute(methodCall.getType());
+        //ensure type params are not included
+        methodCallTypeByArgs = JavaPsiFacade.getElementFactory(method.getProject())
+          .createRawSubstitutor(method).substitute(methodCallTypeByArgs);
+        QuickFixAction.registerQuickFixAction(highlightInfo, 
+                                              getFixRange(methodCall),
+                                              QUICK_FIX_FACTORY.createMethodReturnFix(containerMethod, methodCallTypeByArgs, true));
+      }
+    }
+  }
+
   private static String buildOneLineMismatchDescription(@NotNull PsiExpressionList list,
                                                         @NotNull MethodCandidateInfo candidateInfo,
                                                         @NotNull Ref<PsiElement> elementToHighlight) {
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/afterInferFromArgs.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/afterInferFromArgs.java
new file mode 100644
index 000000000000..ef32b152820d
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/afterInferFromArgs.java
@@ -0,0 +1,12 @@
+import java.util.List;
+
+// "Make 'bar' return 'java.util.List<java.lang.String>'" "true"
+public class Foo {
+  <T> java.util.List<T> foo(T t) {
+    return null;
+  }
+
+  List<String> bar() {
+    return foo("");
+  }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/beforeInferFromArgs.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/beforeInferFromArgs.java
new file mode 100644
index 000000000000..3acffd55e925
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/methodReturn/beforeInferFromArgs.java
@@ -0,0 +1,10 @@
+// "Make 'bar' return 'java.util.List<java.lang.String>'" "true"
+public class Foo {
+  <T> java.util.List<T> foo(T t) {
+    return null;
+  }
+
+  String bar() {
+    return fo<caret>o("");
+  }
+}