IDEA-164822 Migrate to Stream API: support toMap/groupingBy scenarios

java/java-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithCollectFix.java

@@ -17,6 +17,7 @@ package com.intellij.codeInspection.streamMigration;
 
 import com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.InitializerUsageStatus;
 import com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.MapOp;
+import com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.TerminalBlock;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.text.StringUtil;
@@ -24,7 +25,10 @@ import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.JavaCodeStyleManager;
 import com.intellij.psi.codeStyle.SuggestedNameInfo;
 import com.intellij.psi.codeStyle.VariableKind;
+import com.intellij.psi.util.InheritanceUtil;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.PsiUtil;
+import com.siyeh.ig.psiutils.ExpressionUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -60,65 +64,47 @@ class ReplaceWithCollectFix extends MigrateToStreamFix {
   PsiElement migrate(@NotNull Project project,
                      @NotNull PsiLoopStatement loopStatement,
                      @NotNull PsiStatement body,
-                     @NotNull StreamApiMigrationInspection.TerminalBlock tb) {
-    final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project);
-    final PsiMethodCallExpression methodCallExpression = tb.getSingleMethodCall();
-
-    if (methodCallExpression == null) return null;
+                     @NotNull TerminalBlock tb) {
+    PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    if (call == null) return null;
 
     restoreComments(loopStatement, body);
     if (!tb.hasOperations() && StreamApiMigrationInspection.isAddAllCall(tb) && loopStatement instanceof PsiForeachStatement) {
-      PsiExpression iteratedValue = ((PsiForeachStatement)loopStatement).getIteratedValue();
-      if (iteratedValue == null) return null;
-      final PsiType iteratedValueType = iteratedValue.getType();
-      final PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
-      final String qualifierText = qualifierExpression != null ? qualifierExpression.getText() : "";
-      final String collectionText =
-        iteratedValueType instanceof PsiArrayType ? "java.util.Arrays.asList(" + iteratedValue.getText() + ")" :
-        getIteratedValueText(iteratedValue);
-      final String callText = StringUtil.getQualifiedName(qualifierText, "addAll(" + collectionText + ");");
-      return loopStatement.replace(elementFactory.createStatementFromText(callText, loopStatement));
+      return handleAddAll(loopStatement, factory, call);
     }
-    PsiExpression itemToAdd = methodCallExpression.getArgumentList().getExpressions()[0];
-    PsiType addedType = getAddedElementType(methodCallExpression);
+    PsiExpression qualifierExpression = call.getMethodExpression().getQualifierExpression();
+    PsiLocalVariable variable = StreamApiMigrationInspection.extractCollectionVariable(qualifierExpression);
+    if(variable != null && InheritanceUtil.isInheritor(variable.getType(), CommonClassNames.JAVA_UTIL_MAP)) {
+      PsiElement result = handleToMap(loopStatement, tb, call, variable);
+      if (result != null) return result;
+    }
+
+    PsiExpression itemToAdd = call.getArgumentList().getExpressions()[0];
+    PsiType addedType = getAddedElementType(call);
     if (addedType == null) addedType = itemToAdd.getType();
+
+    if(variable == null && qualifierExpression instanceof PsiMethodCallExpression) {
+      PsiMethodCallExpression qualifierCall = (PsiMethodCallExpression)qualifierExpression;
+      if(StreamApiMigrationInspection.isCallOf(qualifierCall, CommonClassNames.JAVA_UTIL_MAP, "computeIfAbsent")) {
+        PsiElement result = handleComputeIfAbsent(loopStatement, tb, itemToAdd, qualifierCall);
+        if (result != null) return result;
+      }
+    }
+
     StringBuilder builder = generateStream(new MapOp(tb.getLastOperation(), itemToAdd, tb.getVariable(), addedType));
 
-    final PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
-    final PsiLocalVariable variable = StreamApiMigrationInspection.extractCollectionVariable(qualifierExpression);
     if (variable != null) {
       InitializerUsageStatus status = StreamApiMigrationInspection.getInitializerUsageStatus(variable, loopStatement);
       if(status != InitializerUsageStatus.UNKNOWN) {
         PsiExpression initializer = variable.getInitializer();
         LOG.assertTrue(initializer != null);
-        PsiMethodCallExpression toArrayExpression =
-          StreamApiMigrationInspection.extractToArrayExpression(loopStatement, methodCallExpression);
-        if(toArrayExpression != null) {
-          PsiType type = initializer.getType();
-          if(type instanceof PsiClassType) {
-            String replacement = StreamApiMigrationInspection.COLLECTION_TO_ARRAY.get(((PsiClassType)type).rawType().getCanonicalText());
-            if(replacement != null) {
-              builder.append(".").append(replacement);
-              PsiExpression[] args = toArrayExpression.getArgumentList().getExpressions();
-              if(args.length == 0) {
-                builder.append("()");
-              } else {
-                if(args.length != 1 || !(args[0] instanceof PsiNewExpression)) return null;
-                PsiNewExpression newArray = (PsiNewExpression)args[0];
-                PsiType arrayType = newArray.getType();
-                if(arrayType == null) return null;
-                String name = arrayType.getCanonicalText();
-                builder.append('(').append(name).append("::new)");
-              }
-              PsiElement result =
-                toArrayExpression.replace(elementFactory.createExpressionFromText(builder.toString(), toArrayExpression));
-              removeLoop(loopStatement);
-              if(status != InitializerUsageStatus.AT_WANTED_PLACE) {
-                variable.delete();
-              }
-              return result;
-            }
+        PsiElement toArrayConversion = handleToArray(builder, initializer, loopStatement, call);
+        if(toArrayConversion != null) {
+          if(status != InitializerUsageStatus.AT_WANTED_PLACE) {
+            variable.delete();
           }
+          return toArrayConversion;
         }
         PsiElement nextStatement = PsiTreeUtil.skipSiblingsForward(loopStatement, PsiComment.class, PsiWhiteSpace.class);
         String comparatorText = StreamApiMigrationInspection.tryExtractSortComparatorText(nextStatement, variable);
@@ -126,29 +112,163 @@ class ReplaceWithCollectFix extends MigrateToStreamFix {
           builder.append(".sorted(").append(comparatorText).append(")");
           nextStatement.delete();
         }
-        String callText = builder.append(".collect(java.util.stream.Collectors.")
+        String callText = builder.append(".collect(" + CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".")
           .append(createInitializerReplacementText(qualifierExpression.getType(), initializer))
           .append(")").toString();
         return replaceInitializer(loopStatement, variable, initializer, callText, status);
       }
     }
-    final String qualifierText = qualifierExpression != null ? qualifierExpression.getText() + "." : "";
+    String qualifierText = qualifierExpression != null ? qualifierExpression.getText() + "." : "";
 
     JavaCodeStyleManager codeStyleManager = JavaCodeStyleManager.getInstance(project);
     SuggestedNameInfo suggestedNameInfo =
       codeStyleManager.suggestVariableName(VariableKind.LOCAL_VARIABLE, null, null, addedType, false);
     if (suggestedNameInfo.names.length == 0) {
-      suggestedNameInfo = codeStyleManager.suggestVariableName(VariableKind.LOCAL_VARIABLE, "item", null, itemToAdd.getType(), false);
+      suggestedNameInfo = codeStyleManager.suggestVariableName(VariableKind.LOCAL_VARIABLE, "item", null, null, false);
     }
-    String varName = codeStyleManager.suggestUniqueVariableName(suggestedNameInfo, methodCallExpression, false).names[0];
+    String varName = codeStyleManager.suggestUniqueVariableName(suggestedNameInfo, call, false).names[0];
 
     PsiExpression forEachBody =
-      elementFactory.createExpressionFromText(qualifierText + "add(" + varName + ")", qualifierExpression);
-    final String callText =
-      builder.append(".forEach(").append(varName).append("->").append(forEachBody.getText()).append(");").toString();
-    return loopStatement.replace(elementFactory.createStatementFromText(callText, loopStatement));
+      factory.createExpressionFromText(qualifierText + "add(" + varName + ")", qualifierExpression);
+    String callText = builder.append(".forEach(").append(varName).append("->").append(forEachBody.getText()).append(");").toString();
+    return loopStatement.replace(factory.createStatementFromText(callText, loopStatement));
   }
 
+  @Nullable
+  private static PsiElement handleComputeIfAbsent(@NotNull PsiLoopStatement loopStatement, @NotNull TerminalBlock tb,
+                                                  PsiExpression itemToAdd, PsiMethodCallExpression qualifierCall) {
+    PsiLocalVariable variable;
+    variable = StreamApiMigrationInspection.extractCollectionVariable(qualifierCall.getMethodExpression().getQualifierExpression());
+    if (variable == null || !InheritanceUtil.isInheritor(variable.getType(), CommonClassNames.JAVA_UTIL_MAP)) return null;
+    InitializerUsageStatus status = StreamApiMigrationInspection.getInitializerUsageStatus(variable, loopStatement);
+    if(status == InitializerUsageStatus.UNKNOWN) return null;
+    PsiExpression[] computeArgs = qualifierCall.getArgumentList().getExpressions();
+    if(!(computeArgs[1] instanceof PsiLambdaExpression)) return null;
+    PsiExpression ctor = LambdaUtil.extractSingleExpressionFromBody(((PsiLambdaExpression)computeArgs[1]).getBody());
+    PsiType mapType = variable.getType();
+    PsiType valueType = PsiUtil.substituteTypeParameter(mapType, CommonClassNames.JAVA_UTIL_MAP, 1, false);
+    if(valueType == null) return null;
+    String downstreamCollector = CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + "." + createInitializerReplacementText(valueType, ctor);
+    if(!ExpressionUtils.isReferenceTo(itemToAdd, tb.getVariable())) {
+      downstreamCollector = CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".mapping(" +
+                            tb.getVariable().getName() + "->" + itemToAdd.getText() + "," + downstreamCollector + ")";
+    }
+    StringBuilder builder = generateStream(tb.getLastOperation());
+    builder.append(".collect(" + CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".groupingBy(")
+      .append(LambdaUtil.createLambda(tb.getVariable(), computeArgs[0]));
+    PsiExpression initializer = variable.getInitializer();
+    LOG.assertTrue(initializer != null);
+    if (!isHashMap(variable)) {
+      builder.append(",()->").append(initializer.getText()).append(",").append(downstreamCollector);
+    }
+    else if (!(CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + "." + "toList()").equals(downstreamCollector)) {
+      builder.append(",").append(downstreamCollector);
+    }
+    builder.append("))");
+    return replaceInitializer(loopStatement, variable, initializer, builder.toString(), status);
+  }
+
+  @Nullable
+  private static PsiElement handleAddAll(@NotNull PsiLoopStatement loopStatement, PsiElementFactory factory, PsiMethodCallExpression call) {
+    PsiExpression iteratedValue = ((PsiForeachStatement)loopStatement).getIteratedValue();
+    if (iteratedValue == null) return null;
+    PsiExpression qualifierExpression = call.getMethodExpression().getQualifierExpression();
+    String qualifierText = qualifierExpression != null ? qualifierExpression.getText() : "";
+    String collectionText =
+      iteratedValue.getType() instanceof PsiArrayType ? CommonClassNames.JAVA_UTIL_ARRAYS + ".asList(" + iteratedValue.getText() + ")" :
+      getIteratedValueText(iteratedValue);
+    String callText = StringUtil.getQualifiedName(qualifierText, "addAll(" + collectionText + ");");
+    return loopStatement.replace(factory.createStatementFromText(callText, loopStatement));
+  }
+
+  @Nullable
+  private static PsiElement handleToMap(@NotNull PsiLoopStatement loopStatement,
+                                        @NotNull TerminalBlock tb,
+                                        PsiMethodCallExpression call,
+                                        PsiLocalVariable variable) {
+    PsiExpression[] args = call.getArgumentList().getExpressions();
+    if(args.length < 2) return null;
+    String methodName = call.getMethodExpression().getReferenceName();
+    if(methodName == null) return null;
+    InitializerUsageStatus status = StreamApiMigrationInspection.getInitializerUsageStatus(variable, loopStatement);
+    if(status == InitializerUsageStatus.UNKNOWN) return null;
+    Project project = loopStatement.getProject();
+    PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
+    PsiExpression merger;
+    JavaCodeStyleManager codeStyleManager = JavaCodeStyleManager.getInstance(project);
+    String aVar = codeStyleManager.suggestUniqueVariableName("a", call, true);
+    String bVar = codeStyleManager.suggestUniqueVariableName("b", call, true);
+    switch(methodName) {
+      case "put":
+        merger = factory.createExpressionFromText("("+aVar+","+bVar+")->"+bVar, call);
+        break;
+      case "putIfAbsent":
+        merger = factory.createExpressionFromText("("+aVar+","+bVar+")->"+aVar, call);
+        break;
+      case "merge":
+        if(args.length != 3) return null;
+        merger = args[2];
+        break;
+      default:
+        return null;
+    }
+    StringBuilder collector = new StringBuilder(CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS+".toMap(");
+    collector.append(LambdaUtil.createLambda(tb.getVariable(), args[0])).append(',')
+      .append(LambdaUtil.createLambda(tb.getVariable(), args[1])).append(',')
+      .append(merger.getText());
+    PsiExpression initializer = variable.getInitializer();
+    LOG.assertTrue(initializer != null);
+    if(!isHashMap(variable)) {
+      collector.append(",()->").append(initializer.getText());
+    }
+    collector.append(")");
+    String callText = generateStream(tb.getLastOperation()).append(".collect(").append(collector).append(")").toString();
+    return replaceInitializer(loopStatement, variable, initializer, callText, status);
+  }
+
+  private static boolean isHashMap(PsiLocalVariable variable) {
+    PsiExpression initializer = variable.getInitializer();
+    LOG.assertTrue(initializer != null);
+    PsiClass initializerClass = PsiUtil.resolveClassInClassTypeOnly(initializer.getType());
+    PsiClass varClass = PsiUtil.resolveClassInClassTypeOnly(variable.getType());
+    return initializerClass != null &&
+           varClass != null &&
+           CommonClassNames.JAVA_UTIL_HASH_MAP.equals(initializerClass.getQualifiedName()) &&
+           CommonClassNames.JAVA_UTIL_MAP.equals(varClass.getQualifiedName());
+  }
+
+  @Nullable
+  private static PsiElement handleToArray(StringBuilder builder,
+                                          PsiExpression initializer,
+                                          PsiLoopStatement loopStatement,
+                                          PsiMethodCallExpression methodCallExpression) {
+    PsiMethodCallExpression toArrayExpression =
+      StreamApiMigrationInspection.extractToArrayExpression(loopStatement, methodCallExpression);
+    if (toArrayExpression == null) return null;
+    PsiType type = initializer.getType();
+    if (!(type instanceof PsiClassType)) return null;
+    String replacement = StreamApiMigrationInspection.COLLECTION_TO_ARRAY.get(((PsiClassType)type).rawType().getCanonicalText());
+    if (replacement == null) return null;
+
+    builder.append(".").append(replacement);
+    PsiExpression[] args = toArrayExpression.getArgumentList().getExpressions();
+    if(args.length == 0) {
+      builder.append("()");
+    } else {
+      if(args.length != 1 || !(args[0] instanceof PsiNewExpression)) return null;
+      PsiNewExpression newArray = (PsiNewExpression)args[0];
+      PsiType arrayType = newArray.getType();
+      if(arrayType == null) return null;
+      String name = arrayType.getCanonicalText();
+      builder.append('(').append(name).append("::new)");
+    }
+    PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(toArrayExpression.getProject());
+    PsiElement result = toArrayExpression.replace(elementFactory.createExpressionFromText(builder.toString(), toArrayExpression));
+    removeLoop(loopStatement);
+    return result;
+  }
+
+  @NotNull
   private static String createInitializerReplacementText(PsiType varType, PsiExpression initializer) {
     final PsiType initializerType = initializer.getType();
     final PsiClassType rawType = initializerType instanceof PsiClassType ? ((PsiClassType)initializerType).rawType() : null;
@@ -164,9 +284,6 @@ class ReplaceWithCollectFix extends MigrateToStreamFix {
               rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_COLLECTION))) {
       return "toSet()";
     }
-    else if (rawType != null) {
-      return "toCollection(" + rawType.getClassName() + "::new)";
-    }
     else {
       return "toCollection(() -> " + initializer.getText() + ")";
     }

java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java

@@ -256,68 +256,100 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
   }
 
   static boolean isAddAllCall(TerminalBlock tb) {
-    final PsiVariable variable = tb.getVariable();
-    final PsiMethodCallExpression methodCallExpression = tb.getSingleMethodCall();
-    LOG.assertTrue(methodCallExpression != null);
-    if (!ExpressionUtils.isReferenceTo(methodCallExpression.getArgumentList().getExpressions()[0], variable)) return false;
-    PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
-    if(qualifierExpression == null || qualifierExpression instanceof PsiThisExpression) {
-      PsiMethod method = PsiTreeUtil.getParentOfType(methodCallExpression, PsiMethod.class);
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    LOG.assertTrue(call != null);
+    if (!ExpressionUtils.isReferenceTo(call.getArgumentList().getExpressions()[0], tb.getVariable())) return false;
+    if (!"add".equals(call.getMethodExpression().getReferenceName())) return false;
+    PsiExpression qualifierExpression = call.getMethodExpression().getQualifierExpression();
+    if (qualifierExpression == null || qualifierExpression instanceof PsiThisExpression) {
+      PsiMethod method = PsiTreeUtil.getParentOfType(call, PsiMethod.class);
       return method == null || !method.getName().equals("addAll");
     }
     return true;
   }
 
+  @Nullable
+  private static PsiClass extractQualifierClass(TerminalBlock tb, PsiMethodCallExpression methodCallExpression) {
+    PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
+    PsiClass qualifierClass = null;
+    if (qualifierExpression instanceof PsiReferenceExpression) {
+      if (ReferencesSearch.search(tb.getVariable(), new LocalSearchScope(qualifierExpression)).findFirst() != null) {
+        return null;
+      }
+      final PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
+      if (resolve instanceof PsiVariable) {
+        if (ReferencesSearch.search(resolve, new LocalSearchScope(methodCallExpression.getArgumentList())).findFirst() != null) {
+          return null;
+        }
+      }
+      qualifierClass = PsiUtil.resolveClassInType(qualifierExpression.getType());
+    }
+    else if (qualifierExpression == null || qualifierExpression instanceof PsiThisExpression) {
+      final PsiClass enclosingClass = PsiTreeUtil.getParentOfType(methodCallExpression, PsiClass.class);
+      if (PsiUtil.getEnclosingStaticElement(methodCallExpression, enclosingClass) == null) {
+        qualifierClass = enclosingClass;
+      }
+    }
+    return qualifierClass;
+  }
+
+  @Contract("null, _, _ -> false")
+  static boolean isCallOf(PsiMethodCallExpression call, String className, String... methodNames) {
+    if (call == null) return false;
+    PsiReferenceExpression methodExpression = call.getMethodExpression();
+    String name = methodExpression.getReferenceName();
+    if (!ArrayUtil.contains(name, methodNames)) return false;
+    PsiMethod maybeMapMethod = call.resolveMethod();
+    if (maybeMapMethod == null ||
+        maybeMapMethod.getParameterList().getParametersCount() != call.getArgumentList().getExpressions().length) {
+      return false;
+    }
+    PsiClass containingClass = maybeMapMethod.getContainingClass();
+    if (containingClass == null) return false;
+    if (className.equals(containingClass.getQualifiedName())) return true;
+    PsiMethod[] superMethods = maybeMapMethod.findDeepestSuperMethods();
+    return StreamEx.of(superMethods).map(PsiMember::getContainingClass).nonNull().map(PsiClass::getQualifiedName).has(className);
+  }
+
+  private static boolean isCollectMapCall(PsiLoopStatement loop, TerminalBlock tb) {
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    if (!isCallOf(call, CommonClassNames.JAVA_UTIL_MAP, "merge", "put", "putIfAbsent")) return false;
+    PsiReferenceExpression methodExpression = call.getMethodExpression();
+    PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
+    return extractQualifierClass(tb, call) != null && !tb.dependsOn(qualifierExpression) && canCollect(loop, call);
+  }
+
   private static boolean isCollectCall(TerminalBlock tb) {
-    final PsiMethodCallExpression methodCallExpression = tb.getSingleMethodCall();
-    if (methodCallExpression != null) {
-      final PsiReferenceExpression methodExpression = methodCallExpression.getMethodExpression();
-      final PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
-      PsiClass qualifierClass = null;
-      if (qualifierExpression instanceof PsiReferenceExpression) {
-        if (ReferencesSearch.search(tb.getVariable(), new LocalSearchScope(qualifierExpression)).findFirst() != null) {
-          return false;
-        }
-        final PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
-        if (resolve instanceof PsiVariable) {
-          if (ReferencesSearch.search(resolve, new LocalSearchScope(methodCallExpression.getArgumentList())).findFirst() != null) {
-            return false;
-          }
-        }
-        qualifierClass = PsiUtil.resolveClassInType(qualifierExpression.getType());
-      }
-      else if (qualifierExpression == null || qualifierExpression instanceof PsiThisExpression) {
-        final PsiClass enclosingClass = PsiTreeUtil.getParentOfType(methodCallExpression, PsiClass.class);
-        if (PsiUtil.getEnclosingStaticElement(methodCallExpression, enclosingClass) == null) {
-          qualifierClass = enclosingClass;
-        }
-      }
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return false;
+    PsiReferenceExpression methodExpression = call.getMethodExpression();
+    PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
 
-      if (qualifierClass != null && 
-          InheritanceUtil.isInheritor(qualifierClass, false, CommonClassNames.JAVA_UTIL_COLLECTION)) {
+    if (tb.dependsOn(qualifierExpression)) return false;
 
-        if (tb.intermediateExpressions().anyMatch(expression -> isExpressionDependsOnUpdatedCollections(expression, qualifierExpression))) {
-          return false;
-        }
-
-        final PsiElement resolve = methodExpression.resolve();
-        if (resolve instanceof PsiMethod &&
-            "add".equals(((PsiMethod)resolve).getName()) &&
-            ((PsiMethod)resolve).getParameterList().getParametersCount() == 1) {
-          final PsiExpression[] args = methodCallExpression.getArgumentList().getExpressions();
-          if (args.length == 1) {
-            if (args[0] instanceof PsiCallExpression) {
-              final PsiMethod method = ((PsiCallExpression)args[0]).resolveMethod();
-              return method != null && !method.hasTypeParameters() && !isThrowsCompatible(method);
-            }
-            return true;
-          }
-        }
+    if (extractQualifierClass(tb, call) != null) return true;
+    if (qualifierExpression instanceof PsiMethodCallExpression) {
+      PsiMethodCallExpression qualifierCall = (PsiMethodCallExpression)qualifierExpression;
+      if (isCallOf(qualifierCall, CommonClassNames.JAVA_UTIL_MAP, "computeIfAbsent")) {
+        PsiExpression[] args = qualifierCall.getArgumentList().getExpressions();
+        if (args.length != 2 || !(args[1] instanceof PsiLambdaExpression)) return false;
+        PsiLambdaExpression lambda = (PsiLambdaExpression)args[1];
+        PsiExpression body = LambdaUtil.extractSingleExpressionFromBody(lambda.getBody());
+        if (!(body instanceof PsiNewExpression)) return false;
+        PsiExpressionList ctorArgs = ((PsiNewExpression)body).getArgumentList();
+        return ctorArgs != null && ctorArgs.getExpressions().length == 0;
       }
     }
     return false;
   }
 
+  @Contract("_, null -> false")
+  static boolean canCollect(PsiLoopStatement statement, PsiMethodCallExpression call) {
+    if(call == null) return false;
+    PsiLocalVariable variable = extractCollectionVariable(call.getMethodExpression().getQualifierExpression());
+    return variable != null && getInitializerUsageStatus(variable, statement) != UNKNOWN;
+  }
+
   private static boolean isExpressionDependsOnUpdatedCollections(PsiExpression condition,
                                                                  PsiExpression qualifierExpression) {
     final PsiElement collection = qualifierExpression instanceof PsiReferenceExpression
@@ -610,9 +642,12 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
           if(addAll) {
             methodName = "addAll";
           } else {
-            PsiMethodCallExpression methodCallExpression = tb.getSingleMethodCall();
-            if(canCollect(statement, methodCallExpression)) {
-              if(extractToArrayExpression(statement, methodCallExpression) != null)
+            PsiMethodCallExpression call = tb.getSingleMethodCall();
+            if(call != null && call.getMethodExpression().getQualifierExpression() instanceof PsiMethodCallExpression) {
+              call = (PsiMethodCallExpression)call.getMethodExpression().getQualifierExpression();
+            }
+            if(canCollect(statement, call)) {
+              if(extractToArrayExpression(statement, call) != null)
                 methodName = "toArray";
               else
                 methodName = "collect";
@@ -623,6 +658,9 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
           }
           registerProblem(statement, methodName, new ReplaceWithCollectFix(methodName));
         }
+        else if (isCollectMapCall(statement, tb) && (REPLACE_TRIVIAL_FOREACH || tb.hasOperations())) {
+          registerProblem(statement, "collect", new ReplaceWithCollectFix("collect"));
+        }
         // do not replace for(T e : arr) {} with Arrays.stream(arr).forEach(e -> {}) even if flag is set
         else if (SUGGEST_FOREACH &&
                  (tb.hasOperations() || (!(source instanceof ArrayStream) && (REPLACE_TRIVIAL_FOREACH || !isTrivial(body, statement))))) {
@@ -673,13 +711,6 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       }
     }
 
-    boolean canCollect(PsiLoopStatement statement, PsiMethodCallExpression methodCallExpression) {
-      if(methodCallExpression == null) return false;
-      PsiLocalVariable variable = extractCollectionVariable(methodCallExpression.getMethodExpression().getQualifierExpression());
-      if(variable == null) return false;
-      return getInitializerUsageStatus(variable, statement) != UNKNOWN;
-    }
-
     void handleSingleReturn(PsiLoopStatement statement, TerminalBlock tb) {
       PsiReturnStatement returnStatement = (PsiReturnStatement)tb.getSingleStatement();
       PsiExpression value = returnStatement.getReturnValue();
@@ -1512,5 +1543,9 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     public static TerminalBlock from(StreamSource source, @NotNull PsiStatement body) {
       return new TerminalBlock(source, source.myVariable, body).extractOperations();
     }
+
+    boolean dependsOn(PsiExpression qualifier) {
+      return intermediateExpressions().anyMatch(expression -> isExpressionDependsOnUpdatedCollections(expression, qualifier));
+    }
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectToMap.java

@@ -0,0 +1,17 @@
+// "Replace with collect" "true"
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+      Map<String, Integer> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.toMap(s -> s, s -> 1, (a, b) -> a + b));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectToMapPut.java

@@ -0,0 +1,17 @@
+// "Replace with collect" "true"
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+      HashMap<String, Integer> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.toMap(s -> s, s -> 1, (a, b) -> b, HashMap::new));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectToMapPutIfAbsent.java

@@ -0,0 +1,17 @@
+// "Replace with collect" "true"
+import java.util.Arrays;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+      Map<String, Integer> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.toMap(s -> s, s -> 1, (a, b) -> a, LinkedHashMap::new));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterGroupingByArrayList.java

@@ -0,0 +1,14 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<Integer, ArrayList<String>> test(String... list) {
+      Map<Integer, ArrayList<String>> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.groupingBy(String::length, Collectors.toCollection(ArrayList::new)));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterGroupingByCustomCollection.java

@@ -0,0 +1,14 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+      Map<Integer, List<String>> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.groupingBy(String::length, Collectors.mapping(String::trim, Collectors.toCollection(LinkedList::new))));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterGroupingByDownstream.java

@@ -0,0 +1,14 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+      HashMap<Integer, List<String>> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.groupingBy(String::length, HashMap::new, Collectors.mapping(String::trim, Collectors.toList())));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterGroupingBySimple.java

@@ -0,0 +1,14 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+      Map<Integer, List<String>> map = Arrays.stream(list).filter(Objects::nonNull).collect(Collectors.groupingBy(String::length));
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectToMap.java

@@ -0,0 +1,19 @@
+// "Replace with collect" "true"
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+    Map<String, Integer> map = new HashMap<>();
+    for(String s : li<caret>st) {
+      if(s != null) {
+        map.merge(s, 1, (a, b) -> a + b);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectToMapPut.java

@@ -0,0 +1,19 @@
+// "Replace with collect" "true"
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+    HashMap<String, Integer> map = new HashMap<>();
+    for(String s : lis<caret>t) {
+      if(s != null) {
+        map.put(s, 1);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectToMapPutIfAbsent.java

@@ -0,0 +1,19 @@
+// "Replace with collect" "true"
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+    Map<String, Integer> map = new LinkedHashMap<>();
+    for(String s : l<caret>ist) {
+      if(s != null) {
+        map.putIfAbsent(s, 1);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectToMapPutInvalid.java

@@ -0,0 +1,19 @@
+// "Replace with collect" "false"
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  private Map<String, Integer> test(String... list) {
+    HashMap<String, Integer> map = new HashMap<>();
+    for(String s : lis<caret>t) {
+      if(s != null) {
+        map.put(s, 1,);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeGroupingByArrayList.java

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+
+public class Main {
+  private Map<Integer, ArrayList<String>> test(String... list) {
+    Map<Integer, ArrayList<String>> map = new HashMap<>();
+    for(String s : li<caret>st) {
+      if(s != null) {
+        map.computeIfAbsent(s.length(), k -> new ArrayList<>()).add(s);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeGroupingByCustomCollection.java

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+    Map<Integer, List<String>> map = new HashMap<>();
+    for(String s : li<caret>st) {
+      if(s != null) {
+        map.computeIfAbsent(s.length(), k -> new LinkedList<>()).add(s.trim());
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeGroupingByDownstream.java

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+    HashMap<Integer, List<String>> map = new HashMap<>();
+    for(String s : li<caret>st) {
+      if(s != null) {
+        map.computeIfAbsent(s.length(), k -> new ArrayList<>()).add(s.trim());
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeGroupingBySimple.java

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+
+public class Main {
+  private Map<Integer, List<String>> test(String... list) {
+    Map<Integer, List<String>> map = new HashMap<>();
+    for(String s : li<caret>st) {
+      if(s != null) {
+        map.computeIfAbsent(s.length(), k -> new ArrayList<>()).add(s);
+      }
+    }
+    return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(new Main().test("a", "bbb", null, "cc", "dd", "eedasfasdfs", "dd"));
+  }
+}