AdjustFunctionContextFix: suggest to change method accepting lambda/methodRef

Fixes cases 1, 2 from IDEA-174219 Add quick-fixes for incompilable Stream API call chains when using primitive streams

java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightVisitorImpl.java

@@ -19,6 +19,7 @@ import com.intellij.codeHighlighting.Pass;
 import com.intellij.codeInsight.daemon.JavaErrorMessages;
 import com.intellij.codeInsight.daemon.impl.*;
 import com.intellij.codeInsight.daemon.impl.analysis.HighlightUtil.Feature;
+import com.intellij.codeInsight.daemon.impl.quickfix.AdjustFunctionContextFix;
 import com.intellij.codeInsight.daemon.impl.quickfix.QuickFixAction;
 import com.intellij.codeInsight.intention.QuickFixFactory;
 import com.intellij.codeInspection.LocalQuickFixAndIntentionActionOnPsiElement;
@@ -288,9 +289,12 @@ public class HighlightVisitorImpl extends JavaElementVisitor implements Highligh
       myHolder.add(AnnotationsHighlightUtil.checkMemberValueType(value, returnType));
     }
 
-    myHolder.add(AnnotationsHighlightUtil.checkValidAnnotationType(method.getReturnType(), method.getReturnTypeElement()));
+    PsiTypeElement returnTypeElement = method.getReturnTypeElement();
+    myHolder.add(AnnotationsHighlightUtil.checkValidAnnotationType(method.getReturnType(), returnTypeElement));
     final PsiClass aClass = method.getContainingClass();
-    myHolder.add(AnnotationsHighlightUtil.checkCyclicMemberType(method.getReturnTypeElement(), aClass));
+    if (returnTypeElement != null && aClass != null) {
+      myHolder.add(AnnotationsHighlightUtil.checkCyclicMemberType(returnTypeElement, aClass));
+    }
     myHolder.add(AnnotationsHighlightUtil.checkClashesWithSuperMethods(method));
 
     if (!myHolder.hasErrorResults() && aClass != null) {
@@ -358,13 +362,21 @@ public class HighlightVisitorImpl extends JavaElementVisitor implements Highligh
       }
       final Map<PsiElement, String> returnErrors = LambdaUtil.checkReturnTypeCompatible(expression, LambdaUtil.getFunctionalInterfaceReturnType(functionalInterfaceType));
       if (parentInferenceErrorMessage != null && (returnErrors == null || !returnErrors.containsValue(parentInferenceErrorMessage))) {
-        myHolder.add(HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(expression).descriptionAndTooltip(parentInferenceErrorMessage).create());
+        HighlightInfo info =
+          HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(expression).descriptionAndTooltip(parentInferenceErrorMessage)
+            .create();
+        if (returnErrors != null) {
+          returnErrors.keySet().forEach(k -> QuickFixAction.registerQuickFixAction(info, AdjustFunctionContextFix.createFix(k)));
+        }
+        myHolder.add(info);
       }
       else if (returnErrors != null) {
         for (Map.Entry<PsiElement, String> entry : returnErrors.entrySet()) {
-          myHolder.add(HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR)
-                         .range(entry.getKey())
-                         .descriptionAndTooltip(entry.getValue()).create());
+          HighlightInfo info = HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR)
+            .range(entry.getKey())
+            .descriptionAndTooltip(entry.getValue()).create();
+          QuickFixAction.registerQuickFixAction(info, AdjustFunctionContextFix.createFix(entry.getKey()));
+          myHolder.add(info);
         }
       }
     }
@@ -1414,8 +1426,10 @@ public class HighlightVisitorImpl extends JavaElementVisitor implements Highligh
     if (!myHolder.hasErrorResults()) {
       final String badReturnTypeMessage = PsiMethodReferenceUtil.checkReturnType(expression, result, functionalInterfaceType);
       if (badReturnTypeMessage != null) {
-        myHolder.add(
-          HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(expression).descriptionAndTooltip(badReturnTypeMessage).create());
+        HighlightInfo info =
+          HighlightInfo.newHighlightInfo(HighlightInfoType.ERROR).range(expression).descriptionAndTooltip(badReturnTypeMessage).create();
+        QuickFixAction.registerQuickFixAction(info, AdjustFunctionContextFix.createFix(expression));
+        myHolder.add(info);
       }
     }
 

java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/AdjustFunctionContextFix.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2000-2017 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 com.intellij.codeInsight.daemon.impl.quickfix;
+
+import com.intellij.codeInsight.daemon.QuickFixBundle;
+import com.intellij.codeInspection.LocalQuickFixAndIntentionActionOnPsiElement;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.*;
+import com.intellij.psi.util.InheritanceUtil;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ObjectUtils;
+import com.siyeh.ig.callMatcher.CallMapper;
+import com.siyeh.ig.callMatcher.CallMatcher;
+import com.siyeh.ig.psiutils.ExpressionUtils;
+import com.siyeh.ig.psiutils.StreamApiUtil;
+import org.jetbrains.annotations.Contract;
+import org.jetbrains.annotations.Nls;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.function.Function;
+
+public class AdjustFunctionContextFix extends LocalQuickFixAndIntentionActionOnPsiElement {
+  private static final Function<PsiMethodCallExpression, Function<PsiType, String>>
+    MAP_NAME_ADJUSTER = (PsiMethodCallExpression call) -> (PsiType type) -> {
+    PsiExpression qualifier = call.getMethodExpression().getQualifierExpression();
+    if (qualifier != null) {
+      PsiType inType = StreamApiUtil.getStreamElementType(qualifier.getType());
+      if (type.equals(inType)) return "map";
+    }
+    if (PsiType.INT.equals(type)) return "mapToInt";
+    if (PsiType.LONG.equals(type)) return "mapToLong";
+    if (PsiType.DOUBLE.equals(type)) return "mapToDouble";
+    return "mapToObj";
+  };
+  private static final Function<PsiType, String> FLAT_MAP_NAME_ADJUSTER = type -> {
+    if(InheritanceUtil.isInheritor(type, CommonClassNames.JAVA_UTIL_STREAM_INT_STREAM)) return "flatMapToInt";
+    if(InheritanceUtil.isInheritor(type, CommonClassNames.JAVA_UTIL_STREAM_LONG_STREAM)) return "flatMapToLong";
+    if(InheritanceUtil.isInheritor(type, CommonClassNames.JAVA_UTIL_STREAM_DOUBLE_STREAM)) return "flatMapToDouble";
+    if(InheritanceUtil.isInheritor(type, CommonClassNames.JAVA_UTIL_STREAM_STREAM)) return "flatMap";
+    return null;
+  };
+
+  private static final CallMapper<Function<PsiType, String>> METHOD_NAME_ADJUSTER = new CallMapper<Function<PsiType, String>>()
+    .register(
+      CallMatcher.anyOf(
+        CallMatcher.instanceCall(CommonClassNames.JAVA_UTIL_STREAM_INT_STREAM, "map", "mapToLong", "mapToDouble"),
+        CallMatcher.instanceCall(CommonClassNames.JAVA_UTIL_STREAM_LONG_STREAM, "map", "mapToInt", "mapToDouble"),
+        CallMatcher.instanceCall(CommonClassNames.JAVA_UTIL_STREAM_DOUBLE_STREAM, "map", "mapToInt", "mapToLong")
+      ), MAP_NAME_ADJUSTER)
+    .register(
+      CallMatcher.instanceCall(CommonClassNames.JAVA_UTIL_STREAM_STREAM, "flatMap", "flatMapToInt", "flatMapToLong", "flatMapToDouble"),
+      FLAT_MAP_NAME_ADJUSTER);
+
+  private final String myOriginalName;
+  private final String myNewName;
+
+  protected AdjustFunctionContextFix(@NotNull PsiMethodCallExpression call, @NotNull String targetMethodName) {
+    super(call);
+    myOriginalName = call.getMethodExpression().getReferenceName();
+    myNewName = targetMethodName;
+  }
+
+  @Override
+  public void invoke(@NotNull Project project,
+                     @NotNull PsiFile file,
+                     @Nullable Editor editor,
+                     @NotNull PsiElement startElement,
+                     @NotNull PsiElement endElement) {
+    PsiMethodCallExpression call = ObjectUtils.tryCast(startElement, PsiMethodCallExpression.class);
+    if (call == null) return;
+    ExpressionUtils.bindCallTo(call, myNewName);
+  }
+
+  @NotNull
+  @Override
+  public String getText() {
+    return QuickFixBundle.message("adjust.method.accepting.functional.expression.fix.text", myOriginalName, myNewName);
+  }
+
+  @Nls
+  @NotNull
+  @Override
+  public String getFamilyName() {
+    return QuickFixBundle.message("adjust.method.accepting.functional.expression.fix.family.name");
+  }
+
+  @Contract("null -> null")
+  @Nullable
+  public static AdjustFunctionContextFix createFix(PsiElement context) {
+    if (!(context instanceof PsiExpression)) return null;
+    PsiExpression expression = (PsiExpression)context;
+    PsiFunctionalExpression fn = PsiTreeUtil.getParentOfType(context, PsiFunctionalExpression.class, false);
+    if (fn == null) return null;
+    PsiExpressionList expressionList = ObjectUtils.tryCast(fn.getParent(), PsiExpressionList.class);
+    if (expressionList == null || expressionList.getExpressions().length != 1) return null;
+    PsiMethodCallExpression call = ObjectUtils.tryCast(expressionList.getParent(), PsiMethodCallExpression.class);
+    Function<PsiType, String> remapper = METHOD_NAME_ADJUSTER.mapFirst(call);
+    if (remapper == null) return null;
+    PsiType actualReturnType;
+    if(expression instanceof PsiMethodReferenceExpression) {
+      PsiMethodReferenceExpression methodRef = (PsiMethodReferenceExpression)expression;
+      actualReturnType = PsiMethodReferenceUtil.getMethodReferenceReturnType(methodRef, methodRef.advancedResolve(true));
+    } else {
+      actualReturnType = PsiResolveHelper.ourGraphGuard.doPreventingRecursion(expression, true, () -> expression.getType());
+    }
+    String targetMethodName = remapper.apply(actualReturnType);
+    if (targetMethodName == null) return null;
+    return new AdjustFunctionContextFix(call, targetMethodName);
+  }
+}

java/java-psi-api/src/com/intellij/psi/PsiMethodReferenceUtil.java

@@ -77,6 +77,55 @@ public class PsiMethodReferenceUtil {
     return isReturnTypeCompatible(expression, result, functionalInterfaceType, null);
   }
 
+  /**
+   * Returns actual return type of method reference (not the expected one)
+   *
+   * @param expression a method reference to get the return type of
+   * @param result the result of method reference resolution
+   * @return an actual method reference return type
+   */
+  public static PsiType getMethodReferenceReturnType(PsiMethodReferenceExpression expression, JavaResolveResult result) {
+    PsiSubstitutor subst = result.getSubstitutor();
+
+    PsiType methodReturnType = null;
+    PsiClass containingClass = null;
+    final PsiElement resolve = result.getElement();
+    if (resolve instanceof PsiMethod) {
+      containingClass = ((PsiMethod)resolve).getContainingClass();
+      methodReturnType = PsiTypesUtil.patchMethodGetClassReturnType(expression, (PsiMethod)resolve);
+      if (methodReturnType == null) {
+        methodReturnType = ((PsiMethod)resolve).getReturnType();
+        if (PsiType.VOID.equals(methodReturnType)) {
+          return methodReturnType;
+        }
+
+        methodReturnType = subst.substitute(methodReturnType);
+      }
+    }
+    else if (resolve instanceof PsiClass) {
+      if (PsiEquivalenceUtil.areElementsEquivalent(resolve, JavaPsiFacade.getElementFactory(expression.getProject()).getArrayClass(PsiUtil.getLanguageLevel(expression)))) {
+        final PsiTypeParameter[] typeParameters = ((PsiClass)resolve).getTypeParameters();
+        if (typeParameters.length == 1) {
+          final PsiType arrayComponentType = subst.substitute(typeParameters[0]);
+          if (arrayComponentType == null) {
+            return null;
+          }
+          methodReturnType = arrayComponentType.createArrayType();
+        }
+      }
+      containingClass = (PsiClass)resolve;
+    }
+
+    if (methodReturnType == null) {
+      if (containingClass == null) {
+        return null;
+      }
+      methodReturnType = JavaPsiFacade.getElementFactory(expression.getProject()).createType(containingClass, subst);
+    }
+
+    return PsiUtil.captureToplevelWildcards(methodReturnType, expression);
+  }
+
   private static boolean isReturnTypeCompatible(PsiMethodReferenceExpression expression,
                                                 JavaResolveResult result,
                                                 PsiType functionalInterfaceType,
@@ -90,45 +139,10 @@ public class PsiMethodReferenceUtil {
         return true;
       }
 
-      PsiSubstitutor subst = result.getSubstitutor();
-
-      PsiType methodReturnType = null;
-      PsiClass containingClass = null;
-      final PsiElement resolve = result.getElement();
-      if (resolve instanceof PsiMethod) {
-        containingClass = ((PsiMethod)resolve).getContainingClass();
-        methodReturnType = PsiTypesUtil.patchMethodGetClassReturnType(expression, (PsiMethod)resolve);
-        if (methodReturnType == null) {
-          methodReturnType = ((PsiMethod)resolve).getReturnType();
-          if (PsiType.VOID.equals(methodReturnType)) {
-            return false;
-          }
-
-          methodReturnType = subst.substitute(methodReturnType);
-        }
+      PsiType methodReturnType = getMethodReferenceReturnType(expression, result);
+      if (methodReturnType == null || PsiType.VOID.equals(methodReturnType)) {
+        return false;
       }
-      else if (resolve instanceof PsiClass) {
-        if (PsiEquivalenceUtil.areElementsEquivalent(resolve, JavaPsiFacade.getElementFactory(expression.getProject()).getArrayClass(PsiUtil.getLanguageLevel(expression)))) {
-          final PsiTypeParameter[] typeParameters = ((PsiClass)resolve).getTypeParameters();
-          if (typeParameters.length == 1) {
-            final PsiType arrayComponentType = subst.substitute(typeParameters[0]);
-            if (arrayComponentType == null) {
-              return false;
-            }
-            methodReturnType = arrayComponentType.createArrayType();
-          }
-        }
-        containingClass = (PsiClass)resolve;
-      }
-
-      if (methodReturnType == null) {
-        if (containingClass == null) {
-          return false;
-        }
-        methodReturnType = JavaPsiFacade.getElementFactory(expression.getProject()).createType(containingClass, subst);
-      }
-
-      methodReturnType = PsiUtil.captureToplevelWildcards(methodReturnType, expression);
 
       if (TypeConversionUtil.isAssignable(interfaceReturnType, methodReturnType)) {
         return true;

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/afterFlatMap.java

@@ -0,0 +1,10 @@
+// "Replace 'flatMap()' with 'flatMapToDouble()'" "true"
+import java.util.*;
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    double[][] data = {{0,1,2},{3,4,5}};
+    Arrays.stream(data).flatMapToDouble(array -> Arrays.stream(array)).forEach(System.out::println);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/afterFlatMapMethodRef.java

@@ -0,0 +1,10 @@
+// "Replace 'flatMap()' with 'flatMapToInt()'" "true"
+import java.util.*;
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    int[][] data = {{0,1,2},{3,4,5}};
+    Arrays.stream(data).flatMapToInt(Arrays::stream).forEach(System.out::println);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/afterIntStreamMap.java

@@ -0,0 +1,8 @@
+// "Replace 'map()' with 'mapToDouble()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    IntStream.range(0, 100).mapToDouble(x -> x/1.0).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/afterIntStreamMapMethodRef.java

@@ -0,0 +1,8 @@
+// "Replace 'map()' with 'mapToObj()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    IntStream.range(0, 100).mapToObj(String::valueOf).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/afterLongStreamMapToInt.java

@@ -0,0 +1,8 @@
+// "Replace 'mapToInt()' with 'map()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    LongStream.range(0, 100).map(x -> x*2).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/beforeFlatMap.java

@@ -0,0 +1,10 @@
+// "Replace 'flatMap()' with 'flatMapToDouble()'" "true"
+import java.util.*;
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    double[][] data = {{0,1,2},{3,4,5}};
+    Arrays.stream(data).flatMap(array -> Arrays.<caret>stream(array)).forEach(System.out::println);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/beforeFlatMapMethodRef.java

@@ -0,0 +1,10 @@
+// "Replace 'flatMap()' with 'flatMapToInt()'" "true"
+import java.util.*;
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    int[][] data = {{0,1,2},{3,4,5}};
+    Arrays.stream(data).flatMap(Arrays<caret>::stream).forEach(System.out::println);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/beforeIntStreamMap.java

@@ -0,0 +1,8 @@
+// "Replace 'map()' with 'mapToDouble()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    IntStream.range(0, 100).map(x -> x/<caret>1.0).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/beforeIntStreamMapMethodRef.java

@@ -0,0 +1,8 @@
+// "Replace 'map()' with 'mapToObj()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    IntStream.range(0, 100).map(String<caret>::valueOf).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext/beforeLongStreamMapToInt.java

@@ -0,0 +1,8 @@
+// "Replace 'mapToInt()' with 'map()'" "true"
+import java.util.stream.*;
+
+class Test {
+  void test() {
+    LongStream.range(0, 100).mapToInt(x -> x<caret>*2).forEach(s -> System.out.println(s));
+  }
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/daemon/quickFix/AdjustFunctionContextFixTest.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2000-2017 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 com.intellij.java.codeInsight.daemon.quickFix;
+
+import com.intellij.codeInsight.daemon.quickFix.LightQuickFixParameterizedTestCase;
+
+public class AdjustFunctionContextFixTest extends LightQuickFixParameterizedTestCase {
+  public void test() {
+    doAllTests();
+  }
+
+  @Override
+  protected String getBasePath() {
+    return "/codeInsight/daemonCodeAnalyzer/quickFix/adjustFunctionContext";
+  }
+
+}

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/callMatcher/CallMapper.java

@@ -18,6 +18,7 @@ package com.siyeh.ig.callMatcher;
 import com.intellij.psi.PsiMethodCallExpression;
 import com.intellij.psi.PsiMethodReferenceExpression;
 import one.util.streamex.StreamEx;
+import org.jetbrains.annotations.Contract;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -60,6 +61,7 @@ public class CallMapper<T> {
     return this;
   }
 
+  @Contract("null -> null")
   public T mapFirst(PsiMethodCallExpression call) {
     if (call == null) return null;
     List<CallHandler<T>> functions = myMap.get(call.getMethodExpression().getReferenceName());

resources-en/src/messages/QuickFixBundle.properties

@@ -317,4 +317,7 @@ java.9.merge.module.statements.fix.family.name=Merge with other ''{0}'' directiv
 java.9.merge.module.statements.fix.name=Merge with other ''{0} {1}'' directive
 
 model.create.constructor.quickfix.message=Create constructor ''{0}''
-model.create.constructor.quickfix.message.family.name=Create constructor
+model.create.constructor.quickfix.message.family.name=Create constructor
+
+adjust.method.accepting.functional.expression.fix.family.name=Adjust method accepting functional expression
+adjust.method.accepting.functional.expression.fix.text=Replace ''{0}()'' with ''{1}()''