diff --git a/java/java-impl/src/com/intellij/codeInspection/streamMigration/CollectMigration.java b/java/java-impl/src/com/intellij/codeInspection/streamMigration/CollectMigration.java
index a6b26b75a17b..1baaf5e17731 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/CollectMigration.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/CollectMigration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * 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.
@@ -21,16 +21,26 @@ import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 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.search.searches.ReferencesSearch;
 import com.intellij.psi.util.InheritanceUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
+import com.siyeh.ig.psiutils.EquivalenceChecker;
 import com.siyeh.ig.psiutils.ExpressionUtils;
+import com.siyeh.ig.psiutils.VariableAccessUtils;
+import one.util.streamex.EntryStream;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiFunction;
+
+import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.isCallOf;
+
 /**
  * @author Tagir Valeev
  */
@@ -42,7 +52,7 @@ class CollectMigration extends BaseStreamApiMigration {
   }
 
   @Nullable
-  PsiType getAddedElementType(PsiMethodCallExpression call) {
+  static PsiType getAddedElementType(PsiMethodCallExpression call) {
     JavaResolveResult resolveResult = call.resolveMethodGenerics();
     PsiMethod method = call.resolveMethod();
     if(method == null) return null;
@@ -57,146 +67,28 @@ class CollectMigration extends BaseStreamApiMigration {
     PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
     PsiMethodCallExpression call = tb.getSingleMethodCall();
     if (call == null) return null;
-
+    CollectTerminal terminal = extractCollectTerminal(tb);
+    if (terminal == null) return null;
+    PsiVariable variable = terminal.getTargetVariable();
+    LOG.assertTrue(variable != null);
+    String stream = tb.generate() + terminal.generateIntermediate() + terminal.generateTerminal();
+    InitializerUsageStatus status = StreamApiMigrationInspection.getInitializerUsageStatus(variable, loopStatement);
+    if (status == InitializerUsageStatus.UNKNOWN) return null;
+    PsiElement toReplace = terminal.getElementToReplace();
     restoreComments(loopStatement, body);
-    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;
+    PsiElement result;
+    if (toReplace != null) {
+      result = toReplace.replace(factory.createExpressionFromText(stream, toReplace));
+      removeLoop(loopStatement);
+      if (status != InitializerUsageStatus.AT_WANTED_PLACE) {
+        variable.delete();
       }
     }
-
-    StringBuilder builder = new StringBuilder(tb.add(new MapOp(itemToAdd, tb.getVariable(), addedType)).generate());
-
-    if (variable != null) {
-      InitializerUsageStatus status = StreamApiMigrationInspection.getInitializerUsageStatus(variable, loopStatement);
-      if(status != InitializerUsageStatus.UNKNOWN) {
-        PsiExpression initializer = variable.getInitializer();
-        LOG.assertTrue(initializer != null);
-        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 = tryExtractSortComparatorText(nextStatement, variable);
-        if(comparatorText != null) {
-          builder.append(".sorted(").append(comparatorText).append(")");
-          nextStatement.delete();
-        }
-        String callText = builder.append(".collect(" + CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".")
-          .append(createInitializerReplacementText(qualifierExpression.getType(), initializer))
-          .append(")").toString();
-        return replaceInitializer(loopStatement, variable, initializer, callText, status);
-      }
+    else {
+      result = replaceInitializer(loopStatement, variable, variable.getInitializer(), stream, status);
     }
-    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, null, false);
-    }
-    String varName = codeStyleManager.suggestUniqueVariableName(suggestedNameInfo, call, false).names[0];
-
-    PsiExpression forEachBody =
-      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 = new StringBuilder(tb.generate());
-    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 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 = tb.generate() + ".collect(" + collector + ")";
-    return replaceInitializer(loopStatement, variable, initializer, callText, status);
+    terminal.cleanUp();
+    return result;
   }
 
   private static boolean isHashMap(PsiLocalVariable variable) {
@@ -211,93 +103,556 @@ class CollectMigration extends BaseStreamApiMigration {
   }
 
   @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;
+  static PsiLocalVariable extractQualifierVariable(TerminalBlock tb, PsiMethodCallExpression call) {
+    PsiExpression qualifierExpression = PsiUtil.skipParenthesizedExprDown(call.getMethodExpression().getQualifierExpression());
+    if (!(qualifierExpression instanceof PsiReferenceExpression)) return null;
+    PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
+    if (!(resolve instanceof PsiLocalVariable)) return null;
+    PsiLocalVariable variable = (PsiLocalVariable)resolve;
+    if (tb.getVariable() != variable && VariableAccessUtils.variableIsUsed(variable, call.getArgumentList())) return null;
+    return variable;
+  }
 
-    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;
+  @Nullable
+  static CollectTerminal extractCollectTerminal(TerminalBlock tb) {
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    if (call == null) return null;
+    PsiReferenceExpression methodExpression = call.getMethodExpression();
+    PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
+    if (tb.dependsOn(qualifierExpression)) return null;
+
+    List<BiFunction<TerminalBlock, PsiMethodCallExpression, CollectTerminal>> extractors =
+      Arrays.asList(AddingTerminal::tryExtract, GroupingTerminal::tryExtract, ToMapTerminal::tryExtract);
+
+    CollectTerminal terminal = StreamEx.of(extractors).map(extractor -> extractor.apply(tb, call)).nonNull().findFirst().orElse(null);
+    if (terminal != null) {
+      PsiVariable variable = terminal.getTargetVariable();
+      if (variable == null || variable.getInitializer() == null) return null;
+      terminal = includePostStatements(terminal, tb.getMainLoop());
+    }
+    return terminal;
+  }
+
+  static CollectTerminal includePostStatements(CollectTerminal terminal, PsiLoopStatement loop) {
+    List<BiFunction<CollectTerminal, PsiElement, CollectTerminal>> wrappers =
+      Arrays.asList(SortingTerminal::tryWrap, (t, e) -> ToArrayTerminal.tryWrap(t, loop, e));
+    PsiElement nextStatement = loop;
+    while (true) {
+      nextStatement = PsiTreeUtil.skipSiblingsForward(nextStatement, PsiComment.class, PsiWhiteSpace.class);
+      CollectTerminal wrapped = null;
+      for (BiFunction<CollectTerminal, PsiElement, CollectTerminal> wrapper : wrappers) {
+        wrapped = wrapper.apply(terminal, nextStatement);
+        if (wrapped != null) {
+          terminal = wrapped;
+          break;
+        }
+      }
+      if (wrapped == null) {
+        return terminal;
+      }
+    }
+  }
+
+  interface CollectTerminal {
+    @Nullable
+    default PsiElement getElementToReplace() { return null; }
+
+    default String getMethodName() { return "collect"; }
+
+    @Nullable
+    PsiVariable getTargetVariable();
+
+    default String generateIntermediate() { return ""; }
+
+    String generateTerminal();
+
+    default void cleanUp() {}
+
+    default boolean isTrivial() {
+      return generateIntermediate().isEmpty();
+    }
+  }
+
+  static class AddingTerminal implements CollectTerminal {
+    private @Nullable PsiVariable myTarget;
+    private final PsiType myTargetType;
+    private final PsiExpression myInitializer;
+    private final PsiVariable myElement;
+    private final PsiMethodCallExpression myAddCall;
+
+    AddingTerminal(@NotNull PsiVariable target,
+                   PsiVariable element,
+                   PsiMethodCallExpression addCall) {
+      this(target.getType(), target.getInitializer(), element, addCall);
+      if (myInitializer instanceof PsiNewExpression) {
+        final PsiExpressionList argumentList = ((PsiNewExpression)myInitializer).getArgumentList();
+        if (argumentList != null && argumentList.getExpressions().length == 0) {
+          myTarget = target;
+        }
+      }
+    }
+
+    AddingTerminal(@NotNull PsiType targetType,
+                   PsiExpression initializer,
+                   PsiVariable element,
+                   PsiMethodCallExpression addCall) {
+      myTarget = null;
+      myTargetType = targetType;
+      myInitializer = initializer;
+      myElement = element;
+      myAddCall = addCall;
+    }
+
+    @Override
+    @Nullable
+    public PsiVariable getTargetVariable() {
+      return myTarget;
+    }
+
+    PsiVariable getElementVariable() {
+      return myElement;
+    }
+
+    PsiExpression getMapping() {
+      return myAddCall.getArgumentList().getExpressions()[0];
+    }
+
+    @Override
+    public String generateIntermediate() {
+      PsiType addedType = getAddedElementType(myAddCall);
+      PsiExpression mapping = getMapping();
+      if (addedType == null) addedType = mapping.getType();
+      return new MapOp(mapping, myElement, addedType).createReplacement();
+    }
+
+    public String generateCollector() {
+      String collector;
+      PsiType initializerType = myInitializer.getType();
+      PsiClassType rawType = initializerType instanceof PsiClassType ? ((PsiClassType)initializerType).rawType() : null;
+      PsiClassType rawVarType = myTargetType instanceof PsiClassType ? ((PsiClassType)myTargetType).rawType() : null;
+      if (rawType != null && rawVarType != null &&
+          rawType.equalsToText(CommonClassNames.JAVA_UTIL_ARRAY_LIST) &&
+          (rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_LIST) || rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_COLLECTION))) {
+        collector = "toList()";
+      }
+      else if (rawType != null && rawVarType != null &&
+               rawType.equalsToText(CommonClassNames.JAVA_UTIL_HASH_SET) &&
+               (rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_SET) ||
+                rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_COLLECTION))) {
+        collector = "toSet()";
+      }
+      else {
+        collector = "toCollection(() -> " + myInitializer.getText() + ")";
+      }
+      return CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + "." + collector;
+    }
+
+    @Override
+    public String generateTerminal() {
+      return ".collect(" + generateCollector() + ")";
+    }
+
+    @Nullable
+    static AddingTerminal tryExtract(TerminalBlock tb, PsiMethodCallExpression call) {
+      if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return null;
+      PsiExpression qualifierExpression = call.getMethodExpression().getQualifierExpression();
+      if (qualifierExpression == null) return null;
+      PsiExpression count = tb.getCountExpression();
+      PsiLocalVariable variable = extractQualifierVariable(tb, call);
+      if (variable != null) {
+        AddingTerminal terminal = new AddingTerminal(variable, tb.getVariable(), call);
+        if (count == null) return terminal;
+        // like "list.add(x); if(list.size() >= limit) break;"
+        if (!(count instanceof PsiMethodCallExpression)) return null;
+        PsiMethodCallExpression sizeCall = (PsiMethodCallExpression)count;
+        PsiExpression sizeQualifier = sizeCall.getMethodExpression().getQualifierExpression();
+        if (isCallOf(sizeCall, CommonClassNames.JAVA_UTIL_COLLECTION, "size") &&
+            EquivalenceChecker.getCanonicalPsiEquivalence().expressionsAreEquivalent(sizeQualifier, qualifierExpression) &&
+            InheritanceUtil.isInheritor(PsiUtil.resolveClassInClassTypeOnly(variable.getType()), CommonClassNames.JAVA_UTIL_LIST)) {
+          return terminal;
+        }
+      }
+      return null;
+    }
+  }
+
+  static class GroupingTerminal implements CollectTerminal {
+    private final AddingTerminal myDownstream;
+    private final PsiLocalVariable myTarget;
+    private final PsiExpression myKeyExpression;
+
+    GroupingTerminal(AddingTerminal downstream, PsiLocalVariable target, PsiExpression expression) {
+      myDownstream = downstream;
+      myTarget = target;
+      myKeyExpression = expression;
+    }
+
+    @Override
+    public boolean isTrivial() {
+      return false;
+    }
+
+    @Nullable
+    @Override
+    public PsiVariable getTargetVariable() {
+      return myTarget;
+    }
+
+    @Override
+    public String generateTerminal() {
+      String downstreamCollector = myDownstream.generateCollector();
+      PsiVariable elementVariable = myDownstream.getElementVariable();
+      if (!ExpressionUtils.isReferenceTo(myDownstream.getMapping(), myDownstream.getElementVariable())) {
+        downstreamCollector = CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".mapping(" +
+                              myDownstream.getElementVariable().getName() + "->" + myDownstream.getMapping().getText() + "," +
+                              downstreamCollector + ")";
+      }
+      StringBuilder builder = new StringBuilder();
+      builder.append(".collect(" + CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".groupingBy(")
+        .append(LambdaUtil.createLambda(elementVariable, myKeyExpression));
+      PsiExpression initializer = myTarget.getInitializer();
+      LOG.assertTrue(initializer != null);
+      if (!isHashMap(myTarget)) {
+        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 builder.toString();
+    }
+
+    @Nullable
+    public static GroupingTerminal tryExtract(TerminalBlock tb, PsiMethodCallExpression call) {
+      if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return null;
+      PsiReferenceExpression methodExpression = call.getMethodExpression();
+      PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
+
+      if (qualifierExpression instanceof PsiMethodCallExpression && tb.getCountExpression() == null) {
+        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 null;
+          PsiLambdaExpression lambda = (PsiLambdaExpression)args[1];
+          PsiExpression body = LambdaUtil.extractSingleExpressionFromBody(lambda.getBody());
+          if (!(body instanceof PsiNewExpression)) return null;
+          PsiExpressionList ctorArgs = ((PsiNewExpression)body).getArgumentList();
+          if (ctorArgs != null && ctorArgs.getExpressions().length == 0) {
+            PsiLocalVariable variable = extractQualifierVariable(tb, qualifierCall);
+            if (variable != null && variable.getInitializer() instanceof PsiNewExpression) {
+              PsiType mapType = variable.getType();
+              PsiType valueType = PsiUtil.substituteTypeParameter(mapType, CommonClassNames.JAVA_UTIL_MAP, 1, false);
+              if (valueType == null) return null;
+              AddingTerminal adding = new AddingTerminal(valueType, body, tb.getVariable(), call);
+              return new GroupingTerminal(adding, variable, args[0]);
+            }
+          }
+        }
+      }
+      return null;
+    }
+  }
+
+  static class ToMapTerminal implements CollectTerminal {
+    private final PsiMethodCallExpression myMapUpdateCall;
+    private final PsiLocalVariable myTargetVariable;
+    private final PsiVariable myElementVariable;
+
+    ToMapTerminal(PsiMethodCallExpression call, PsiVariable elementVariable, PsiLocalVariable variable) {
+      myMapUpdateCall = call;
+      myTargetVariable = variable;
+      myElementVariable = elementVariable;
+    }
+
+    @Nullable
+    @Override
+    public PsiVariable getTargetVariable() {
+      return myTargetVariable;
+    }
+
+    @Override
+    public String generateTerminal() {
+      PsiExpression[] args = myMapUpdateCall.getArgumentList().getExpressions();
+      LOG.assertTrue(args.length >= 2);
+      String methodName = myMapUpdateCall.getMethodExpression().getReferenceName();
+      LOG.assertTrue(methodName != null);
+      Project project = myMapUpdateCall.getProject();
+      String merger;
+      JavaCodeStyleManager codeStyleManager = JavaCodeStyleManager.getInstance(project);
+      String aVar = codeStyleManager.suggestUniqueVariableName("a", myMapUpdateCall, true);
+      String bVar = codeStyleManager.suggestUniqueVariableName("b", myMapUpdateCall, true);
+      switch (methodName) {
+        case "put":
+          merger = "(" + aVar + "," + bVar + ")->" + bVar;
+          break;
+        case "putIfAbsent":
+          merger = "(" + aVar + "," + bVar + ")->" + aVar;
+          break;
+        case "merge":
+          LOG.assertTrue(args.length == 3);
+          merger = args[2].getText();
+          break;
+        default:
+          return null;
+      }
+      StringBuilder collector = new StringBuilder(".collect(" + CommonClassNames.JAVA_UTIL_STREAM_COLLECTORS + ".toMap(");
+      collector.append(LambdaUtil.createLambda(myElementVariable, args[0])).append(',')
+        .append(LambdaUtil.createLambda(myElementVariable, args[1])).append(',')
+        .append(merger);
+      PsiExpression initializer = myTargetVariable.getInitializer();
+      LOG.assertTrue(initializer != null);
+      if (!isHashMap(myTargetVariable)) {
+        collector.append(",()->").append(initializer.getText());
+      }
+      collector.append("))");
+      return collector.toString();
+    }
+
+    @Nullable
+    static ToMapTerminal tryExtract(TerminalBlock tb, PsiMethodCallExpression call) {
+      if (tb.getCountExpression() != null ||
+          !isCallOf(call, CommonClassNames.JAVA_UTIL_MAP, "merge", "put", "putIfAbsent")) {
+        return null;
+      }
+      PsiLocalVariable variable = extractQualifierVariable(tb, call);
+      if (variable == null) return null;
+      PsiExpression initializer = variable.getInitializer();
+      if (!(initializer instanceof PsiNewExpression)) return null;
+      PsiExpressionList argumentList = ((PsiNewExpression)initializer).getArgumentList();
+      if (argumentList == null || argumentList.getExpressions().length != 0) return null;
+      return new ToMapTerminal(call, tb.getVariable(), variable);
+    }
+  }
+
+  static class SortingTerminal implements CollectTerminal {
+    private final CollectTerminal myDownstream;
+    private final PsiExpression myComparator;
+    private final PsiStatement myStatement;
+
+    SortingTerminal(CollectTerminal downstream, PsiStatement statement, PsiExpression comparator) {
+      myDownstream = downstream;
+      myStatement = statement;
+      myComparator = comparator;
+    }
+
+    @Override
+    public String getMethodName() {
+      return myDownstream.getMethodName();
+    }
+
+    @Nullable
+    @Override
+    public PsiVariable getTargetVariable() {
+      return myDownstream.getTargetVariable();
+    }
+
+    @Override
+    public String generateIntermediate() {
+      return myDownstream.generateIntermediate() + ".sorted("
+             + (myComparator == null ? "" : myComparator.getText()) + ")";
+    }
+
+    @Override
+    public String generateTerminal() {
+      return myDownstream.generateTerminal();
+    }
+
+    @Override
+    public void cleanUp() {
+      myDownstream.cleanUp();
+      myStatement.delete();
+    }
+
+    @Contract("null, _ -> null")
+    public static CollectTerminal tryWrap(CollectTerminal terminal, PsiElement element) {
+      PsiVariable list = terminal.getTargetVariable();
+      if (list == null || !(element instanceof PsiExpressionStatement)) return null;
+      PsiExpression expression = ((PsiExpressionStatement)element).getExpression();
+      if (!(expression instanceof PsiMethodCallExpression)) return null;
+      PsiMethodCallExpression methodCall = (PsiMethodCallExpression)expression;
+      PsiReferenceExpression methodExpression = methodCall.getMethodExpression();
+      if (!"sort".equals(methodExpression.getReferenceName())) return null;
+      PsiMethod method = methodCall.resolveMethod();
+      if (method == null) return null;
+      PsiClass containingClass = method.getContainingClass();
+      if (containingClass == null) return null;
+      PsiExpression listExpression = null;
+      PsiExpression comparatorExpression = null;
+      if (CommonClassNames.JAVA_UTIL_COLLECTIONS.equals(containingClass.getQualifiedName())) {
+        PsiExpression[] args = methodCall.getArgumentList().getExpressions();
+        if (args.length == 1) {
+          listExpression = args[0];
+        }
+        else if (args.length == 2) {
+          listExpression = args[0];
+          comparatorExpression = args[1];
+        }
+        else {
+          return null;
+        }
+      }
+      else if (InheritanceUtil.isInheritor(containingClass, CommonClassNames.JAVA_UTIL_LIST)) {
+        listExpression = methodExpression.getQualifierExpression();
+        PsiExpression[] args = methodCall.getArgumentList().getExpressions();
+        if (args.length != 1) return null;
+        comparatorExpression = args[0];
+      }
+      if (!(listExpression instanceof PsiReferenceExpression) || !((PsiReferenceExpression)listExpression).isReferenceTo(list)) {
+        return null;
+      }
+      if (ExpressionUtils.isNullLiteral(comparatorExpression)) {
+        comparatorExpression = null;
+      }
+      return new SortingTerminal(terminal, (PsiExpressionStatement)element, comparatorExpression);
+    }
+  }
+
+  static class ToArrayTerminal implements CollectTerminal {
+    static final Map<String, String> INTERMEDIATE_STEPS = EntryStream.of(
+      CommonClassNames.JAVA_UTIL_ARRAY_LIST, "",
+      "java.util.LinkedList", "",
+      CommonClassNames.JAVA_UTIL_HASH_SET, ".distinct()",
+      "java.util.LinkedHashSet", ".distinct()",
+      "java.util.TreeSet", ".distinct().sorted()"
+    ).toMap();
+
+    private final CollectTerminal myUpstream;
+    private final String myIntermediate;
+    private final PsiMethodCallExpression myToArrayExpression;
+    private final String mySupplier;
+
+    ToArrayTerminal(CollectTerminal upstream,
+                    String intermediate,
+                    String supplier,
+                    PsiMethodCallExpression toArrayExpression) {
+      myUpstream = upstream;
+      mySupplier = supplier;
+      myIntermediate = intermediate;
+      myToArrayExpression = toArrayExpression;
+    }
+
+    @Override
+    public boolean isTrivial() {
+      return false;
+    }
+
+    @Override
+    public String getMethodName() {
+      return "toArray";
+    }
+
+    @Nullable
+    @Override
+    public PsiElement getElementToReplace() {
+      return myToArrayExpression;
+    }
+
+    @Nullable
+    @Override
+    public PsiVariable getTargetVariable() {
+      return myUpstream.getTargetVariable();
+    }
+
+    @Override
+    public String generateIntermediate() {
+      return myUpstream.generateIntermediate() + myIntermediate;
+    }
+
+    @Override
+    public String generateTerminal() {
+      return ".toArray(" + mySupplier + ")";
+    }
+
+    @Override
+    public void cleanUp() {
+      myUpstream.cleanUp();
+    }
+
+    @Contract("null, _, _ -> null")
+    public static ToArrayTerminal tryWrap(CollectTerminal terminal, PsiLoopStatement loopStatement, PsiElement element) {
+      PsiVariable collectionVariable = terminal.getTargetVariable();
+      if (collectionVariable == null || StreamApiMigrationInspection.getInitializerUsageStatus(collectionVariable, loopStatement)
+                                        == InitializerUsageStatus.UNKNOWN) {
+        return null;
+      }
+      PsiExpression initializer = collectionVariable.getInitializer();
+      if (initializer == null) return null;
+      PsiClass aClass = PsiUtil.resolveClassInClassTypeOnly(initializer.getType());
+      if (aClass == null) return null;
+      String intermediateSteps = INTERMEDIATE_STEPS.get(aClass.getQualifiedName());
+      if (intermediateSteps == null) return null;
+
+      PsiMethodCallExpression toArrayExpression =
+        extractToArrayExpression(element, loopStatement, collectionVariable);
+      if (toArrayExpression == null) return null;
+      PsiExpression[] args = toArrayExpression.getArgumentList().getExpressions();
+      String supplier;
+      if (args.length == 0) {
+        supplier = "";
+      }
+      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();
+        supplier = name + "::new";
+      }
+      return new ToArrayTerminal(terminal, intermediateSteps, supplier, toArrayExpression);
+    }
+
+    @Nullable
+    static PsiMethodCallExpression extractToArrayExpression(PsiElement nextElement, PsiLoopStatement statement, PsiVariable collectionVariable) {
+      // return collection.toArray() or collection.toArray(new Type[0]) or collection.toArray(new Type[collection.size()]);
+      PsiExpression toArrayCandidate;
+      if (nextElement instanceof PsiReturnStatement) {
+        toArrayCandidate = ((PsiReturnStatement)nextElement).getReturnValue();
+      }
+      else {
+        PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(nextElement);
+        if (assignment != null) {
+          toArrayCandidate = assignment.getRExpression();
+        }
+        else if (nextElement instanceof PsiDeclarationStatement) {
+          PsiElement[] elements = ((PsiDeclarationStatement)nextElement).getDeclaredElements();
+          if (elements.length == 1 && elements[0] instanceof PsiLocalVariable) {
+            toArrayCandidate = ((PsiLocalVariable)elements[0]).getInitializer();
+          }
+          else {
+            return null;
+          }
+        }
+        else {
+          return null;
+        }
+      }
+      if (!(toArrayCandidate instanceof PsiMethodCallExpression)) return null;
+      PsiMethodCallExpression call = (PsiMethodCallExpression)toArrayCandidate;
+      PsiReferenceExpression methodExpression = call.getMethodExpression();
+      if (!"toArray".equals(methodExpression.getReferenceName())) return null;
+      if (!ExpressionUtils.isReferenceTo(methodExpression.getQualifierExpression(), collectionVariable)) return null;
+
+      if (!(nextElement instanceof PsiReturnStatement) && !ReferencesSearch.search(collectionVariable)
+        .forEach(ref ->
+                   ref.getElement() == collectionVariable || PsiTreeUtil.isAncestor(statement, ref.getElement(), false) ||
+                   PsiTreeUtil.isAncestor(toArrayCandidate, ref.getElement(), false)
+        )) {
+        return null;
+      }
+
+      PsiExpression[] args = call.getArgumentList().getExpressions();
+      if (args.length == 0) return call;
+      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)");
+      PsiExpression[] dimensions = newArray.getArrayDimensions();
+      if (dimensions.length != 1) return null;
+      if (ExpressionUtils.isLiteral(dimensions[0], 0)) return call;
+      if (!(dimensions[0] instanceof PsiMethodCallExpression)) return null;
+      PsiMethodCallExpression maybeSizeCall = (PsiMethodCallExpression)dimensions[0];
+      if (!isCallOf(maybeSizeCall, CommonClassNames.JAVA_UTIL_COLLECTION, "size")) return null;
+      PsiExpression sizeQualifier = maybeSizeCall.getMethodExpression().getQualifierExpression();
+      if (!ExpressionUtils.isReferenceTo(sizeQualifier, collectionVariable)) return null;
+      return call;
     }
-    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;
-    final PsiClassType rawVarType = varType instanceof PsiClassType ? ((PsiClassType)varType).rawType() : null;
-    if (rawType != null && rawVarType != null &&
-        rawType.equalsToText(CommonClassNames.JAVA_UTIL_ARRAY_LIST) &&
-        (rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_LIST) || rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_COLLECTION))) {
-      return "toList()";
-    }
-    else if (rawType != null && rawVarType != null &&
-             rawType.equalsToText(CommonClassNames.JAVA_UTIL_HASH_SET) &&
-             (rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_SET) ||
-              rawVarType.equalsToText(CommonClassNames.JAVA_UTIL_COLLECTION))) {
-      return "toSet()";
-    }
-    else {
-      return "toCollection(() -> " + initializer.getText() + ")";
-    }
-  }
-
-  /**
-   *
-   * @param element sort statement candidate (must be PsiExpressionStatement)
-   * @param list list which should be sorted
-   * @return comparator string representation, empty string if natural order is used or null if given statement is not sort statement
-   */
-  @Contract(value = "null, _ -> null")
-  private static String tryExtractSortComparatorText(PsiElement element, PsiVariable list) {
-    if(!(element instanceof PsiExpressionStatement)) return null;
-    PsiExpression expression = ((PsiExpressionStatement)element).getExpression();
-    if(!(expression instanceof PsiMethodCallExpression)) return null;
-    PsiMethodCallExpression methodCall = (PsiMethodCallExpression)expression;
-    PsiReferenceExpression methodExpression = methodCall.getMethodExpression();
-    if(!"sort".equals(methodExpression.getReferenceName())) return null;
-    PsiMethod method = methodCall.resolveMethod();
-    if(method == null) return null;
-    PsiClass containingClass = method.getContainingClass();
-    if(containingClass == null) return null;
-    PsiExpression listExpression = null;
-    PsiExpression comparatorExpression = null;
-    if(CommonClassNames.JAVA_UTIL_COLLECTIONS.equals(containingClass.getQualifiedName())) {
-      PsiExpression[] args = methodCall.getArgumentList().getExpressions();
-      if(args.length == 1) {
-        listExpression = args[0];
-      } else if(args.length == 2) {
-        listExpression = args[0];
-        comparatorExpression = args[1];
-      } else return null;
-    } else if(InheritanceUtil.isInheritor(containingClass, CommonClassNames.JAVA_UTIL_LIST)) {
-      listExpression = methodExpression.getQualifierExpression();
-      PsiExpression[] args = methodCall.getArgumentList().getExpressions();
-      if(args.length != 1) return null;
-      comparatorExpression = args[0];
-    }
-    if(!(listExpression instanceof PsiReferenceExpression) || !((PsiReferenceExpression)listExpression).isReferenceTo(list)) return null;
-    if(comparatorExpression == null || ExpressionUtils.isNullLiteral(comparatorExpression)) return "";
-    return comparatorExpression.getText();
   }
 }
diff --git a/java/java-impl/src/com/intellij/codeInspection/streamMigration/ForEachMigration.java b/java/java-impl/src/com/intellij/codeInspection/streamMigration/ForEachMigration.java
index ead4cd74c883..0e13a43971f4 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/ForEachMigration.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/ForEachMigration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * 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.
@@ -15,11 +15,20 @@
  */
 package com.intellij.codeInspection.streamMigration;
 
+import com.intellij.codeInspection.LambdaCanBeMethodReferenceInspection;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
+import com.intellij.psi.codeStyle.JavaCodeStyleManager;
+import com.intellij.psi.codeStyle.SuggestedNameInfo;
+import com.intellij.psi.codeStyle.VariableKind;
+import com.siyeh.ig.psiutils.ExpressionUtils;
+import com.siyeh.ig.psiutils.VariableAccessUtils;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import static com.intellij.codeInspection.streamMigration.CollectMigration.getAddedElementType;
 
 /**
  * @author Tagir Valeev
@@ -31,18 +40,55 @@ class ForEachMigration extends BaseStreamApiMigration {
     super(forEachMethodName);
   }
 
+  @Nullable
+  static PsiExpression tryExtractMapExpression(TerminalBlock tb) {
+    PsiMethodCallExpression call = tb.getSingleMethodCall();
+    if(call == null) return null;
+    PsiExpression[] args = call.getArgumentList().getExpressions();
+    if(args.length != 1) return null;
+    PsiExpression arg = args[0];
+    if(ExpressionUtils.isReferenceTo(arg, tb.getVariable())) return null;
+    PsiExpression qualifier = call.getMethodExpression().getQualifierExpression();
+    if(tb.dependsOn(qualifier) ||
+       VariableAccessUtils.variableIsUsed(tb.getVariable(), qualifier) ||
+       StreamApiMigrationInspection.isExpressionDependsOnUpdatedCollections(arg, qualifier) ||
+       !LambdaCanBeMethodReferenceInspection.checkQualifier(qualifier)) return null;
+    return arg;
+  }
+
   @Override
   PsiElement migrate(@NotNull Project project, @NotNull PsiStatement body, @NotNull TerminalBlock tb) {
     PsiLoopStatement loopStatement = tb.getMainLoop();
     restoreComments(loopStatement, body);
 
-    final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project);
+    PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
 
-    String stream = tb.generate(true)+"."+getReplacement()+"(";
-    PsiElement block = tb.convertToElement(elementFactory);
+    PsiExpression mapExpression = tryExtractMapExpression(tb);
+
+    if(mapExpression != null) {
+      PsiMethodCallExpression call = tb.getSingleMethodCall();
+      LOG.assertTrue(call != null);
+      PsiType addedType = getAddedElementType(call);
+      if (addedType == null) addedType = call.getType();
+      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, null, false);
+      }
+      String varName = codeStyleManager.suggestUniqueVariableName(suggestedNameInfo, call, false).names[0];
+
+      String streamText = tb.add(new StreamApiMigrationInspection.MapOp(mapExpression, tb.getVariable(), addedType)).generate();
+      String forEachBody = varName + "->" + call.getMethodExpression().getText() + "(" + varName + ")";
+      String callText = streamText + "." + getReplacement() + "(" + forEachBody + ");";
+      return loopStatement.replace(factory.createStatementFromText(callText, loopStatement));
+    }
+
+    String stream = tb.generate(true) + "." + getReplacement() + "(";
+    PsiElement block = tb.convertToElement(factory);
 
     final String functionalExpressionText = tb.getVariable().getName() + " -> " + wrapInBlock(block);
-    PsiExpressionStatement callStatement = (PsiExpressionStatement)elementFactory
+    PsiExpressionStatement callStatement = (PsiExpressionStatement)factory
       .createStatementFromText(stream + functionalExpressionText + ");", loopStatement);
     callStatement = (PsiExpressionStatement)loopStatement.replace(callStatement);
 
@@ -54,7 +100,7 @@ class ForEachMigration extends BaseStreamApiMigration {
     if (expressions[0] instanceof PsiFunctionalExpression &&
         ((PsiFunctionalExpression)expressions[0]).getFunctionalInterfaceType() == null) {
       callStatement =
-        (PsiExpressionStatement)callStatement.replace(elementFactory.createStatementFromText(
+        (PsiExpressionStatement)callStatement.replace(factory.createStatementFromText(
           stream + "(" + tb.getVariable().getText() + ") -> " + wrapInBlock(block) + ");", callStatement));
     }
     return callStatement;
diff --git a/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java b/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java
index f2ab82dc0f18..4ba40b052dfc 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2016 JetBrains s.r.o.
+ * 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.
@@ -16,7 +16,6 @@
 package com.intellij.codeInspection.streamMigration;
 
 import com.intellij.codeInsight.ExceptionUtil;
-import com.intellij.codeInsight.PsiEquivalenceUtil;
 import com.intellij.codeInsight.daemon.GroupNames;
 import com.intellij.codeInsight.daemon.impl.analysis.HighlightControlFlowUtil;
 import com.intellij.codeInspection.BaseJavaBatchLocalInspectionTool;
@@ -40,9 +39,7 @@ import com.intellij.psi.util.PsiUtil;
 import com.intellij.psi.util.TypeConversionUtil;
 import com.intellij.refactoring.util.RefactoringUtil;
 import com.intellij.util.ArrayUtil;
-import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.ig.psiutils.*;
-import one.util.streamex.EntryStream;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.Nls;
@@ -53,6 +50,7 @@ import javax.swing.*;
 import java.util.*;
 
 import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.InitializerUsageStatus.*;
+import static com.intellij.util.ObjectUtils.tryCast;
 
 /**
  * User: anna
@@ -60,14 +58,6 @@ import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInsp
 public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTool {
   private static final Logger LOG = Logger.getInstance("#" + StreamApiMigrationInspection.class.getName());
 
-  static final Map<String, String> COLLECTION_TO_ARRAY = EntryStream.of(
-    CommonClassNames.JAVA_UTIL_ARRAY_LIST, "toArray",
-    "java.util.LinkedList", "toArray",
-    CommonClassNames.JAVA_UTIL_HASH_SET, "distinct().toArray",
-    "java.util.LinkedHashSet", "distinct().toArray",
-    "java.util.TreeSet", "distinct().sorted().toArray"
-  ).toMap();
-
   public boolean REPLACE_TRIVIAL_FOREACH;
   public boolean SUGGEST_FOREACH;
 
@@ -168,11 +158,10 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
 
   @Nullable
   static PsiVariable extractAccumulator(PsiAssignmentExpression assignment) {
-    if(!(assignment.getLExpression() instanceof PsiReferenceExpression)) return null;
-    PsiReferenceExpression lExpr = (PsiReferenceExpression)assignment.getLExpression();
-    PsiElement accumulator = lExpr.resolve();
-    if(!(accumulator instanceof PsiVariable)) return null;
-    PsiVariable var = (PsiVariable)accumulator;
+    PsiReferenceExpression lExpr = tryCast(assignment.getLExpression(), PsiReferenceExpression.class);
+    if(lExpr == null) return null;
+    PsiVariable var = tryCast(lExpr.resolve(), PsiVariable.class);
+    if(var == null) return null;
     if(JavaTokenType.PLUSEQ.equals(assignment.getOperationTokenType())) {
       return var;
     } else if(JavaTokenType.EQ.equals(assignment.getOperationTokenType())) {
@@ -221,13 +210,13 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     if(variables.size() != 1) return null;
 
     // have single expression which is either ++x or x++ or x+=1 or x=x+1
-    PsiExpression operand = extractIncrementedLValue(expression);
-    if(!(operand instanceof PsiReferenceExpression)) return null;
-    PsiElement element = ((PsiReferenceExpression)operand).resolve();
+    PsiReferenceExpression operand = tryCast(extractIncrementedLValue(expression), PsiReferenceExpression.class);
+    if(operand == null) return null;
+    PsiLocalVariable variable = tryCast(operand.resolve(), PsiLocalVariable.class);
 
     // the referred variable is the same as non-final variable and not used in intermediate operations
-    if (element instanceof PsiLocalVariable && variables.contains(element) && !tb.isReferencedInOperations((PsiVariable)element)) {
-      return (PsiLocalVariable)element;
+    if (variable != null && variables.contains(variable) && !tb.isReferencedInOperations(variable)) {
+      return variable;
     }
     return null;
   }
@@ -266,30 +255,6 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     return !(qualifierExpression instanceof PsiMethodCallExpression);
   }
 
-  @Nullable
-  private static PsiClass extractQualifierClass(TerminalBlock tb, PsiMethodCallExpression methodCallExpression) {
-    PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
-    PsiClass qualifierClass = null;
-    if (qualifierExpression instanceof PsiReferenceExpression) {
-      if (VariableAccessUtils.variableIsUsed(tb.getVariable(), qualifierExpression)) {
-        return null;
-      }
-      final PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
-      if (resolve instanceof PsiVariable &&
-          VariableAccessUtils.variableIsUsed((PsiVariable)resolve, methodCallExpression.getArgumentList())) {
-        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;
@@ -308,14 +273,6 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     return StreamEx.of(superMethods).map(PsiMember::getContainingClass).nonNull().map(PsiClass::getQualifiedName).has(className);
   }
 
-  private static boolean isCollectMapCall(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(tb.getMainLoop(), call);
-  }
-
   private static boolean isCountOperation(List<PsiVariable> nonFinalVariables, TerminalBlock tb) {
     PsiLocalVariable variable = getIncrementedVariable(tb.getSingleExpression(PsiExpression.class), tb, nonFinalVariables);
     PsiExpression counter = tb.getCountExpression();
@@ -333,61 +290,14 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
            getInitializerUsageStatus(variable, tb.getMainLoop()) != UNKNOWN;
   }
 
-  private static boolean isCollectCall(TerminalBlock tb) {
-    PsiMethodCallExpression call = tb.getSingleMethodCall();
-    if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return false;
-    PsiReferenceExpression methodExpression = call.getMethodExpression();
-    PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
-
-    if (tb.dependsOn(qualifierExpression)) return false;
-
-    PsiExpression count = tb.getCountExpression();
-    PsiClass qualifierClass = extractQualifierClass(tb, call);
-    if (qualifierClass != null) {
-      if (count == null) return true;
-      // like "list.add(x); if(list.size() >= limit) break;"
-      if(!(count instanceof PsiMethodCallExpression)) return false;
-      PsiMethodCallExpression sizeCall = (PsiMethodCallExpression)count;
-      PsiExpression sizeQualifier = sizeCall.getMethodExpression().getQualifierExpression();
-      return isCallOf(sizeCall, CommonClassNames.JAVA_UTIL_COLLECTION, "size") &&
-             EquivalenceChecker.getCanonicalPsiEquivalence().expressionsAreEquivalent(sizeQualifier, qualifierExpression) &&
-             InheritanceUtil.isInheritor(qualifierClass, CommonClassNames.JAVA_UTIL_LIST);
-    }
-    if (qualifierExpression instanceof PsiMethodCallExpression && count == null) {
-      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 && extractQualifierClass(tb, qualifierCall) != null;
-      }
-    }
-    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;
-  }
-
-  @Contract("_, null -> false")
-  private static boolean isTrivial(PsiStatement body, PsiLoopStatement loopStatement) {
-    if(!(loopStatement instanceof PsiForeachStatement)) return false;
-    PsiParameter parameter = ((PsiForeachStatement)loopStatement).getIterationParameter();
-    //method reference
+  private static boolean isTrivial(TerminalBlock tb) {
+    PsiVariable variable = tb.getVariable();
     final PsiExpression candidate = LambdaCanBeMethodReferenceInspection
-      .canBeMethodReferenceProblem(body instanceof PsiBlockStatement ? ((PsiBlockStatement)body).getCodeBlock() : body,
-                                   new PsiParameter[]{parameter}, createDefaultConsumerType(parameter.getProject(), parameter), null);
-    if (!(candidate instanceof PsiCallExpression)) {
-      return true;
-    }
+      .canBeMethodReferenceProblem(tb.getSingleStatement(),
+                                   new PsiVariable[]{variable}, createDefaultConsumerType(variable.getProject(), variable), null);
+    if (!(candidate instanceof PsiCallExpression)) return true;
     final PsiMethod method = ((PsiCallExpression)candidate).resolveMethod();
-    return method != null && isThrowsCompatible(method);
+    return method == null;
   }
 
   static boolean isSupported(PsiType type) {
@@ -397,10 +307,6 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     return true;
   }
 
-  private static boolean isThrowsCompatible(PsiMethod method) {
-    return ContainerUtil.find(method.getThrowsList().getReferencedTypes(), type -> !ExceptionUtil.isUncheckedException(type)) != null;
-  }
-
   @Nullable
   private static PsiClassType createDefaultConsumerType(Project project, PsiVariable variable) {
     final JavaPsiFacade psiFacade = JavaPsiFacade.getInstance(project);
@@ -410,34 +316,16 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
 
   static boolean isVariableSuitableForStream(PsiVariable variable, PsiStatement statement, TerminalBlock tb) {
     if(ReferencesSearch.search(variable).forEach(ref -> {
-      PsiElement element = ref.getElement();
-      return !(element instanceof PsiExpression) ||
-             !PsiUtil.isAccessedForWriting((PsiExpression)element) ||
-             tb.operations().anyMatch(op -> op.isWriteAllowed(variable, (PsiExpression)element));
+      PsiExpression expression = tryCast(ref.getElement(), PsiExpression.class);
+      return expression == null ||
+             !PsiUtil.isAccessedForWriting(expression) ||
+             tb.operations().anyMatch(op -> op.isWriteAllowed(variable, expression));
     })) {
       return true;
     }
     return HighlightControlFlowUtil.isEffectivelyFinal(variable, statement, null);
   }
 
-  @Contract("null -> null")
-  static PsiLocalVariable extractCollectionVariable(PsiExpression qualifierExpression) {
-    if (qualifierExpression instanceof PsiReferenceExpression) {
-      final PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
-      if (resolve instanceof PsiLocalVariable) {
-        PsiLocalVariable var = (PsiLocalVariable)resolve;
-        final PsiExpression initializer = var.getInitializer();
-        if (initializer instanceof PsiNewExpression) {
-          final PsiExpressionList argumentList = ((PsiNewExpression)initializer).getArgumentList();
-          if (argumentList != null && argumentList.getExpressions().length == 0) {
-            return var;
-          }
-        }
-      }
-    }
-    return null;
-  }
-
   /**
    * Checks whether variable can be referenced between start and loop entry. Back-edges are also considered, so the actual place
    * where it referenced might be outside of (start, loop entry) interval.
@@ -553,6 +441,46 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     return mapOp == null ? "" : "."+mapOp+"("+variable.getName()+" -> "+variable.getName()+")";
   }
 
+  static boolean isExpressionDependsOnUpdatedCollections(PsiExpression condition,
+                                                         PsiExpression qualifierExpression) {
+    final PsiElement collection = qualifierExpression instanceof PsiReferenceExpression
+                                  ? ((PsiReferenceExpression)qualifierExpression).resolve()
+                                  : null;
+    if (collection != null) {
+      return collection instanceof PsiVariable && VariableAccessUtils.variableIsUsed((PsiVariable)collection, condition);
+    }
+
+    final boolean[] dependsOnCollection = {false};
+    condition.accept(new JavaRecursiveElementWalkingVisitor() {
+      @Override
+      public void visitMethodCallExpression(PsiMethodCallExpression expression) {
+        super.visitMethodCallExpression(expression);
+        final PsiExpression callQualifier = expression.getMethodExpression().getQualifierExpression();
+        if (callQualifier == null ||
+            callQualifier instanceof PsiThisExpression && ((PsiThisExpression)callQualifier).getQualifier() == null ||
+            callQualifier instanceof PsiSuperExpression && ((PsiSuperExpression)callQualifier).getQualifier() == null) {
+          dependsOnCollection[0] = true;
+        }
+      }
+
+      @Override
+      public void visitThisExpression(PsiThisExpression expression) {
+        super.visitThisExpression(expression);
+        if (expression.getQualifier() == null && expression.getParent() instanceof PsiExpressionList) {
+          dependsOnCollection[0] = true;
+        }
+      }
+
+      @Override
+      public void visitClass(PsiClass aClass) {}
+
+      @Override
+      public void visitLambdaExpression(PsiLambdaExpression expression) {}
+    });
+
+    return dependsOnCollection[0];
+  }
+
   private class StreamApiMigrationVisitor extends JavaElementVisitor {
     private final ProblemsHolder myHolder;
     private final boolean myIsOnTheFly;
@@ -621,22 +549,27 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       if (isCountOperation(nonFinalVariables, tb)) {
         return new CountMigration();
       }
-      if (isCollectCall(tb) && nonFinalVariables.isEmpty()) {
-        return findCollectMigration(loop, tb);
+      if (nonFinalVariables.isEmpty()) {
+        CollectMigration.CollectTerminal terminal = CollectMigration.extractCollectTerminal(tb);
+        if(terminal != null && getInitializerUsageStatus(terminal.getTargetVariable(), loop) != UNKNOWN) {
+          boolean addAll = loop instanceof PsiForeachStatement && !tb.hasOperations() && isAddAllCall(tb);
+          // Don't suggest to convert the loop which can be trivially replaced via addAll:
+          // this is covered by UseBulkOperationInspection and ManualArrayToCollectionCopyInspection
+          if(addAll) return null;
+          if(!REPLACE_TRIVIAL_FOREACH && !tb.hasOperations() && terminal.isTrivial()) return null;
+          return new CollectMigration(terminal.getMethodName());
+        }
       }
       if (tb.getCountExpression() != null || tb.isEmpty()) return null;
       if (getAccumulatedVariable(tb, nonFinalVariables) != null) {
         return new SumMigration();
       }
-      if (isCollectMapCall(tb) && nonFinalVariables.isEmpty() && (REPLACE_TRIVIAL_FOREACH || tb.hasOperations())) {
-        return new CollectMigration("collect");
-      }
       Collection<PsiStatement> exitPoints = tb.findExitPoints(controlFlow);
       if (exitPoints == null) return null;
-      // do not replace for(T e : arr) {} with Arrays.stream(arr).forEach(e -> {}) even if flag is set
-      if (SUGGEST_FOREACH && exitPoints.isEmpty() && nonFinalVariables.isEmpty() &&
-          (tb.hasOperations() ||
-           (!(tb.getLastOperation() instanceof ArrayStream) && (REPLACE_TRIVIAL_FOREACH || !isTrivial(body, loop))))) {
+      if (SUGGEST_FOREACH && exitPoints.isEmpty() && nonFinalVariables.isEmpty()) {
+        boolean nonTrivial = tb.hasOperations() || ForEachMigration.tryExtractMapExpression(tb) != null || !isTrivial(tb);
+        // do not replace for(T e : arr) {} with Arrays.stream(arr).forEach(e -> {}) even if REPLACE_TRIVIAL_FOREACH is set
+        if (!nonTrivial && (!REPLACE_TRIVIAL_FOREACH || tb.getLastOperation() instanceof ArrayStream)) return null;
         return new ForEachMigration("forEach");
       }
       if (!tb.hasOperations() && !REPLACE_TRIVIAL_FOREACH) return null;
@@ -672,41 +605,18 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       if (nonFinalVariables.size() == 1) {
         PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(statement);
         if(assignment == null) return null;
-        PsiExpression lValue = assignment.getLExpression();
-        if (!(lValue instanceof PsiReferenceExpression)) return null;
-        PsiElement var = ((PsiReferenceExpression)lValue).resolve();
-        if(!(var instanceof PsiVariable) || !nonFinalVariables.contains(var)) return null;
+        PsiReferenceExpression lValue = tryCast(assignment.getLExpression(), PsiReferenceExpression.class);
+        if (lValue == null) return null;
+        PsiVariable var = tryCast(lValue.resolve(), PsiVariable.class);
+        if(var == null || !nonFinalVariables.contains(var)) return null;
         PsiExpression rValue = assignment.getRExpression();
-        if(rValue == null || VariableAccessUtils.variableIsUsed((PsiVariable)var, rValue)) return null;
+        if(rValue == null || VariableAccessUtils.variableIsUsed(var, rValue)) return null;
         if(tb.getVariable().getType() instanceof PsiPrimitiveType && !ExpressionUtils.isReferenceTo(rValue, tb.getVariable())) return null;
         return new FindFirstMigration();
       }
       return null;
     }
 
-    @Nullable
-    private BaseStreamApiMigration findCollectMigration(PsiLoopStatement statement, TerminalBlock tb) {
-      boolean addAll = statement instanceof PsiForeachStatement && !tb.hasOperations() && isAddAllCall(tb);
-      // Don't suggest to convert the loop which can be trivially replaced via addAll:
-      // this is covered by UseBulkOperationInspection and ManualArrayToCollectionCopyInspection
-      if(addAll) return null;
-      PsiMethodCallExpression call = tb.getSingleMethodCall();
-      if(call != null && call.getMethodExpression().getQualifierExpression() instanceof PsiMethodCallExpression) {
-        call = (PsiMethodCallExpression)call.getMethodExpression().getQualifierExpression();
-      }
-      String methodName;
-      if(canCollect(statement, call)) {
-        if(extractToArrayExpression(statement, call) != null)
-          methodName = "toArray";
-        else
-          methodName = "collect";
-      } else {
-        if (!SUGGEST_FOREACH || tb.getCountExpression() != null) return null;
-        methodName = "forEach";
-      }
-      return new CollectMigration(methodName);
-    }
-
     @Nullable
     private BaseStreamApiMigration findMigrationForReturn(PsiLoopStatement statement, TerminalBlock tb) {
       PsiReturnStatement returnStatement = (PsiReturnStatement)tb.getSingleStatement();
@@ -776,75 +686,6 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     }
   }
 
-  @Nullable
-  static PsiMethodCallExpression extractToArrayExpression(PsiLoopStatement statement, PsiMethodCallExpression expression) {
-    // return collection.toArray() or collection.toArray(new Type[0]) or collection.toArray(new Type[collection.size()]);
-    PsiElement nextElement = PsiTreeUtil.skipSiblingsForward(statement, PsiComment.class, PsiWhiteSpace.class);
-    PsiExpression toArrayCandidate;
-    if (nextElement instanceof PsiReturnStatement) {
-      toArrayCandidate = ((PsiReturnStatement)nextElement).getReturnValue();
-    }
-    else {
-      PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(nextElement);
-      if (assignment != null) {
-        toArrayCandidate = assignment.getRExpression();
-      }
-      else if (nextElement instanceof PsiDeclarationStatement) {
-        PsiElement[] elements = ((PsiDeclarationStatement)nextElement).getDeclaredElements();
-        if (elements.length == 1 && elements[0] instanceof PsiLocalVariable) {
-          toArrayCandidate = ((PsiLocalVariable)elements[0]).getInitializer();
-        }
-        else {
-          return null;
-        }
-      }
-      else {
-        return null;
-      }
-    }
-    if (!(toArrayCandidate instanceof PsiMethodCallExpression)) return null;
-    PsiMethodCallExpression call = (PsiMethodCallExpression)toArrayCandidate;
-    PsiReferenceExpression methodExpression = call.getMethodExpression();
-    if (!"toArray".equals(methodExpression.getReferenceName())) return null;
-    PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
-    if (!(qualifierExpression instanceof PsiReferenceExpression)) return null;
-    PsiLocalVariable collectionVariable = extractCollectionVariable(expression.getMethodExpression().getQualifierExpression());
-    if (collectionVariable == null || !((PsiReferenceExpression)qualifierExpression).isReferenceTo(collectionVariable)) return null;
-    PsiExpression initializer = collectionVariable.getInitializer();
-    if (initializer == null) return null;
-    PsiType type = initializer.getType();
-    if (!(type instanceof PsiClassType) || !COLLECTION_TO_ARRAY.containsKey(((PsiClassType)type).rawType().getCanonicalText())) {
-      return null;
-    }
-
-    if (!(nextElement instanceof PsiReturnStatement) && !ReferencesSearch.search(collectionVariable)
-      .forEach(ref ->
-                 ref.getElement() == collectionVariable || PsiTreeUtil.isAncestor(statement, ref.getElement(), false) ||
-                 PsiTreeUtil.isAncestor(toArrayCandidate, ref.getElement(), false)
-      )) {
-      return null;
-    }
-
-    PsiExpression[] args = call.getArgumentList().getExpressions();
-    if (args.length == 0) return call;
-    if (args.length != 1 || !(args[0] instanceof PsiNewExpression)) return null;
-    PsiNewExpression newArray = (PsiNewExpression)args[0];
-    PsiExpression[] dimensions = newArray.getArrayDimensions();
-    if (dimensions.length != 1) return null;
-    if (ExpressionUtils.isLiteral(dimensions[0], 0)) return call;
-    if (!(dimensions[0] instanceof PsiMethodCallExpression)) return null;
-    PsiMethodCallExpression maybeSizeCall = (PsiMethodCallExpression)dimensions[0];
-    if (maybeSizeCall.getArgumentList().getExpressions().length != 0) return null;
-    PsiReferenceExpression maybeSizeExpression = maybeSizeCall.getMethodExpression();
-    PsiExpression sizeQualifier = maybeSizeExpression.getQualifierExpression();
-    if (sizeQualifier != null &&
-        !("size".equals(maybeSizeExpression.getReferenceName()) &&
-          PsiEquivalenceUtil.areElementsEquivalent(qualifierExpression, sizeQualifier))) {
-      return null;
-    }
-    return call;
-  }
-
   /**
    * Intermediate stream operation representation
    */
@@ -1137,28 +978,24 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     @Nullable
     public static BufferedReaderLines from(PsiWhileStatement whileLoop) {
       // while ((line = br.readLine()) != null)
-      PsiExpression condition = PsiUtil.skipParenthesizedExprDown(whileLoop.getCondition());
-      if(!(condition instanceof PsiBinaryExpression)) return null;
-      PsiBinaryExpression binOp = (PsiBinaryExpression)condition;
+      PsiBinaryExpression binOp = tryCast(PsiUtil.skipParenthesizedExprDown(whileLoop.getCondition()), PsiBinaryExpression.class);
+      if(binOp == null) return null;
       if(!JavaTokenType.NE.equals(binOp.getOperationTokenType())) return null;
       PsiExpression operand = ExpressionUtils.getValueComparedWithNull(binOp);
       if(operand == null) return null;
       PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(PsiUtil.skipParenthesizedExprDown(operand));
       if(assignment == null) return null;
-      PsiExpression lValue = assignment.getLExpression();
-      if(!(lValue instanceof PsiReferenceExpression)) return null;
-      PsiElement element = ((PsiReferenceExpression)lValue).resolve();
-      if(!(element instanceof PsiLocalVariable)) return null;
-      PsiLocalVariable var = (PsiLocalVariable)element;
+      PsiReferenceExpression lValue = tryCast(assignment.getLExpression(), PsiReferenceExpression.class);
+      if(lValue == null) return null;
+      PsiLocalVariable var = tryCast(lValue.resolve(), PsiLocalVariable.class);
+      if(var == null) return null;
       if(!ReferencesSearch.search(var).forEach(ref -> {
         return PsiTreeUtil.isAncestor(whileLoop, ref.getElement(), true);
       })) {
         return null;
       }
-      PsiExpression rValue = PsiUtil.skipParenthesizedExprDown(assignment.getRExpression());
-      if(!(rValue instanceof PsiMethodCallExpression)) return null;
-      PsiMethodCallExpression call = (PsiMethodCallExpression)rValue;
-      if(call.getArgumentList().getExpressions().length != 0) return null;
+      PsiMethodCallExpression call = tryCast(PsiUtil.skipParenthesizedExprDown(assignment.getRExpression()), PsiMethodCallExpression.class);
+      if (call == null || call.getArgumentList().getExpressions().length != 0) return null;
       if(!"readLine".equals(call.getMethodExpression().getReferenceName())) return null;
       PsiExpression readerExpression = call.getMethodExpression().getQualifierExpression();
       if(readerExpression == null) return null;
@@ -1185,17 +1022,14 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       PsiExpression iteratedValue = statement.getIteratedValue();
       if (iteratedValue == null) return null;
 
-      PsiType iteratedValueType = iteratedValue.getType();
+      PsiArrayType iteratedValueType = tryCast(iteratedValue.getType(), PsiArrayType.class);
       PsiParameter parameter = statement.getIterationParameter();
 
-      if (!(iteratedValueType instanceof PsiArrayType) ||
-          !isSupported(((PsiArrayType)iteratedValueType).getComponentType()) ||
-          ((parameter.getType() instanceof PsiPrimitiveType) &&
-           !parameter.getType().equals(((PsiArrayType)iteratedValueType).getComponentType()))) {
-        return null;
+      if (iteratedValueType != null && isSupported(iteratedValueType.getComponentType()) &&
+          (!(parameter.getType() instanceof PsiPrimitiveType) || parameter.getType().equals(iteratedValueType.getComponentType()))) {
+        return new ArrayStream(statement, parameter, iteratedValue);
       }
-
-      return new ArrayStream(statement, parameter, iteratedValue);
+      return null;
     }
   }
 
@@ -1212,8 +1046,9 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
 
     @Contract("null, _ -> false")
     static boolean isRawSubstitution(PsiType iteratedValueType, PsiClass collectionClass) {
-      return iteratedValueType instanceof PsiClassType && PsiUtil
-        .isRawSubstitutor(collectionClass, TypeConversionUtil.getSuperClassSubstitutor(collectionClass, (PsiClassType)iteratedValueType));
+      return iteratedValueType instanceof PsiClassType &&
+             PsiUtil.isRawSubstitutor(collectionClass,
+                                      TypeConversionUtil.getSuperClassSubstitutor(collectionClass, (PsiClassType)iteratedValueType));
     }
 
     @Nullable
@@ -1275,12 +1110,10 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     @Nullable
     public static CountingLoop from(PsiForStatement forStatement) {
       // check that initialization is for(int/long i = <initial_value>;...;...)
-      if(!(forStatement.getInitialization() instanceof PsiDeclarationStatement)) return null;
-      PsiDeclarationStatement initialization = (PsiDeclarationStatement)forStatement.getInitialization();
-      if(initialization.getDeclaredElements().length != 1) return null;
-      PsiElement declaration = initialization.getDeclaredElements()[0];
-      if(!(declaration instanceof PsiLocalVariable)) return null;
-      PsiLocalVariable counter = (PsiLocalVariable)declaration;
+      PsiDeclarationStatement initialization = tryCast(forStatement.getInitialization(), PsiDeclarationStatement.class);
+      if (initialization == null || initialization.getDeclaredElements().length != 1) return null;
+      PsiLocalVariable counter = tryCast(initialization.getDeclaredElements()[0], PsiLocalVariable.class);
+      if(counter == null) return null;
       if(!counter.getType().equals(PsiType.INT) && !counter.getType().equals(PsiType.LONG)) return null;
 
       PsiExpression initializer = counter.getInitializer();
@@ -1290,8 +1123,8 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
       if(!VariableAccessUtils.variableIsIncremented(counter, forStatement.getUpdate())) return null;
 
       // check that condition is like for(...;i<bound;...) or for(...;i<=bound;...)
-      if(!(forStatement.getCondition() instanceof PsiBinaryExpression)) return null;
-      PsiBinaryExpression condition = (PsiBinaryExpression)forStatement.getCondition();
+      PsiBinaryExpression condition = tryCast(forStatement.getCondition(), PsiBinaryExpression.class);
+      if(condition == null) return null;
       IElementType type = condition.getOperationTokenType();
       boolean closed = false;
       PsiExpression bound;
@@ -1311,9 +1144,9 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         bound = condition.getLOperand();
         ref = condition.getROperand();
       } else return null;
-      if(bound == null || !(ref instanceof PsiReferenceExpression) || !((PsiReferenceExpression)ref).isReferenceTo(counter)) return null;
+      if(bound == null || !ExpressionUtils.isReferenceTo(ref, counter)) return null;
       if(!TypeConversionUtil.areTypesAssignmentCompatible(counter.getType(), bound)) return null;
-      return new CountingLoop(forStatement,  counter, initializer, bound, closed);
+      return new CountingLoop(forStatement, counter, initializer, bound, closed);
     }
   }
 }
diff --git a/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java b/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java
index 43c2b0943e71..6e23e89f17e6 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java
@@ -23,7 +23,6 @@ import com.intellij.psi.search.LocalSearchScope;
 import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.ArrayUtil;
-import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.IntArrayList;
 import com.siyeh.ig.psiutils.*;
 import one.util.streamex.StreamEx;
@@ -36,6 +35,7 @@ import java.util.NoSuchElementException;
 import java.util.Objects;
 
 import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.*;
+import static com.intellij.util.ObjectUtils.tryCast;
 
 /**
  * This immutable class represents the code which should be performed
@@ -97,7 +97,7 @@ class TerminalBlock {
     PsiStatement statement = getSingleStatement();
     if(statement instanceof PsiExpressionStatement) {
       PsiExpression expression = ((PsiExpressionStatement)statement).getExpression();
-      return ObjectUtils.tryCast(expression, wantedType);
+      return tryCast(expression, wantedType);
     }
     return null;
   }
@@ -195,16 +195,13 @@ class TerminalBlock {
         PsiDeclarationStatement decl = (PsiDeclarationStatement)first;
         PsiElement[] elements = decl.getDeclaredElements();
         if(elements.length == 1) {
-          PsiElement element = elements[0];
-          if(element instanceof PsiLocalVariable) {
-            PsiLocalVariable declaredVar = (PsiLocalVariable)element;
-            if (isSupported(declaredVar.getType())) {
-              PsiExpression initializer = declaredVar.getInitializer();
-              PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
-              if (initializer != null && ReferencesSearch.search(myVariable, new LocalSearchScope(leftOver)).findFirst() == null) {
-                MapOp op = new MapOp(initializer, myVariable, declaredVar.getType());
-                return new TerminalBlock(this, op, declaredVar, leftOver);
-              }
+          PsiLocalVariable declaredVar = tryCast(elements[0], PsiLocalVariable.class);
+          if (declaredVar != null && isSupported(declaredVar.getType())) {
+            PsiExpression initializer = declaredVar.getInitializer();
+            PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
+            if (initializer != null && ReferencesSearch.search(myVariable, new LocalSearchScope(leftOver)).findFirst() == null) {
+              MapOp op = new MapOp(initializer, myVariable, declaredVar.getType());
+              return new TerminalBlock(this, op, declaredVar, leftOver);
             }
           }
         }
@@ -244,11 +241,9 @@ class TerminalBlock {
     }
     if (tb == null || !ControlFlowUtils.statementBreaksLoop(tb.getSingleStatement(), getMainLoop())) return this;
     FilterOp filter = tb.getLastOperation(FilterOp.class);
-    if(filter == null) return this;
-    PsiExpression condition = PsiUtil.skipParenthesizedExprDown(filter.getExpression());
-    if(!(condition instanceof PsiBinaryExpression)) return this;
-    PsiBinaryExpression binOp = (PsiBinaryExpression)condition;
-    if(!ComparisonUtils.isComparison(binOp)) return this;
+    if (filter == null) return this;
+    PsiBinaryExpression binOp = tryCast(PsiUtil.skipParenthesizedExprDown(filter.getExpression()), PsiBinaryExpression.class);
+    if (binOp == null || !ComparisonUtils.isComparison(binOp)) return this;
     String comparison = filter.isNegated() ? ComparisonUtils.getNegatedComparison(binOp.getOperationTokenType())
                         : binOp.getOperationSign().getText();
     boolean flipped = false;
@@ -276,10 +271,8 @@ class TerminalBlock {
     PsiLocalVariable var = null;
     if (dedicatedCounter) {
       if (!(incrementedValue instanceof PsiReferenceExpression)) return this;
-      PsiElement element = ((PsiReferenceExpression)incrementedValue).resolve();
-      if (!(element instanceof PsiLocalVariable)) return this;
-      var = (PsiLocalVariable)element;
-      if (!ExpressionUtils.isZero(var.getInitializer()) || ReferencesSearch.search(var).findAll().size() != 1) return this;
+      var = tryCast(((PsiReferenceExpression)incrementedValue).resolve(), PsiLocalVariable.class);
+      if (var == null || !ExpressionUtils.isZero(var.getInitializer()) || ReferencesSearch.search(var).findAll().size() != 1) return this;
     }
     PsiExpression limit = flipped ? binOp.getLOperand() : binOp.getROperand();
     if(!ExpressionUtils.isSimpleExpression(limit) || VariableAccessUtils.variableIsUsed(myVariable, limit)) return this;
@@ -308,11 +301,19 @@ class TerminalBlock {
     if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return null;
     PsiExpression[] args = call.getArgumentList().getExpressions();
     if (args.length != 1 || !ExpressionUtils.isReferenceTo(args[0], myVariable)) return null;
-    PsiLocalVariable collectionVariable = extractCollectionVariable(call.getMethodExpression().getQualifierExpression());
-    if (collectionVariable == null || getInitializerUsageStatus(collectionVariable, getMainLoop()) == InitializerUsageStatus.UNKNOWN) {
+    PsiReferenceExpression qualifier = tryCast(call.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class);
+    if (qualifier == null) return null;
+    PsiLocalVariable var = tryCast(qualifier.resolve(), PsiLocalVariable.class);
+    if (var == null) return null;
+    PsiNewExpression initializer = tryCast(var.getInitializer(), PsiNewExpression.class);
+    if (initializer == null) return null;
+    PsiExpressionList argumentList = initializer.getArgumentList();
+    if (argumentList == null ||
+        argumentList.getExpressions().length != 0 ||
+        getInitializerUsageStatus(var, getMainLoop()) == InitializerUsageStatus.UNKNOWN) {
       return null;
     }
-    return collectionVariable;
+    return var;
   }
 
   @NotNull
@@ -367,7 +368,7 @@ class TerminalBlock {
 
   @Nullable
   <T extends Operation> T getLastOperation(Class<T> clazz) {
-    return ObjectUtils.tryCast(getLastOperation(), clazz);
+    return tryCast(getLastOperation(), clazz);
   }
 
   @Nullable
@@ -464,44 +465,4 @@ class TerminalBlock {
   boolean isReferencedInOperations(PsiVariable variable) {
     return intermediateAndSourceExpressions().anyMatch(expr -> VariableAccessUtils.variableIsUsed(variable, expr));
   }
-
-  private static boolean isExpressionDependsOnUpdatedCollections(PsiExpression condition,
-                                                                 PsiExpression qualifierExpression) {
-    final PsiElement collection = qualifierExpression instanceof PsiReferenceExpression
-                                  ? ((PsiReferenceExpression)qualifierExpression).resolve()
-                                  : null;
-    if (collection != null) {
-      return collection instanceof PsiVariable && VariableAccessUtils.variableIsUsed((PsiVariable)collection, condition);
-    }
-
-    final boolean[] dependsOnCollection = {false};
-    condition.accept(new JavaRecursiveElementWalkingVisitor() {
-      @Override
-      public void visitMethodCallExpression(PsiMethodCallExpression expression) {
-        super.visitMethodCallExpression(expression);
-        final PsiExpression callQualifier = expression.getMethodExpression().getQualifierExpression();
-        if (callQualifier == null ||
-            callQualifier instanceof PsiThisExpression && ((PsiThisExpression)callQualifier).getQualifier() == null ||
-            callQualifier instanceof PsiSuperExpression && ((PsiSuperExpression)callQualifier).getQualifier() == null) {
-          dependsOnCollection[0] = true;
-        }
-      }
-
-      @Override
-      public void visitThisExpression(PsiThisExpression expression) {
-        super.visitThisExpression(expression);
-        if (expression.getQualifier() == null && expression.getParent() instanceof PsiExpressionList) {
-          dependsOnCollection[0] = true;
-        }
-      }
-
-      @Override
-      public void visitClass(PsiClass aClass) {}
-
-      @Override
-      public void visitLambdaExpression(PsiLambdaExpression expression) {}
-    });
-
-    return dependsOnCollection[0];
-  }
 }
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectHashSetFieldInitializer.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectHashSetFieldInitializer.java
index c85018fbfda9..149f8a24ed53 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectHashSetFieldInitializer.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectHashSetFieldInitializer.java
@@ -8,8 +8,8 @@ public class Collect {
     }
   }
 
-  Set<String> names = new HashSet<>();
+  final Set<String> names = new HashSet<>();
   void collectNames(List<Person> persons){
-      persons.stream().map(Person::getName).forEach(s -> names.add(s));
+      persons.stream().map(Person::getName).forEach(names::add);
   }
 }
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterSortedToArray.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterSortedToArray.java
new file mode 100644
index 000000000000..c2f0d942321c
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterSortedToArray.java
@@ -0,0 +1,8 @@
+// "Replace with toArray" "true"
+import java.util.*;
+
+public class Test {
+  String[] test(List<String> list) {
+      return list.stream().filter(Objects::nonNull).sorted(String.CASE_INSENSITIVE_ORDER).toArray(String[]::new);
+  }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectHashSetFieldInitializer.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectHashSetFieldInitializer.java
index 99d260c5f082..d3f43f58047e 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectHashSetFieldInitializer.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectHashSetFieldInitializer.java
@@ -8,7 +8,7 @@ public class Collect {
     }
   }
 
-  Set<String> names = new HashSet<>();
+  final Set<String> names = new HashSet<>();
   void collectNames(List<Person> persons){
     for (Person person : pers<caret>ons) {
       names.add(person.getName());
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeDisableForArrays.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeDisableForArrays.java
index 312383341141..e2f89ad96fae 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeDisableForArrays.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeDisableForArrays.java
@@ -4,7 +4,7 @@ import java.util.*;
 class A {
   void fun(List<String>... lists) {
     for (List<String> list : li<caret>sts) {
-      System.out.println(list);
+      list.add("");
     }
   }
 
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeSortedToArray.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeSortedToArray.java
new file mode 100644
index 000000000000..80e69f2e219c
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeSortedToArray.java
@@ -0,0 +1,15 @@
+// "Replace with toArray" "true"
+import java.util.*;
+
+public class Test {
+  String[] test(List<String> list) {
+    List<String> result = new LinkedList<>();
+    for(String str : li<caret>st) {
+      if(str != null) {
+        result.add(str);
+      }
+    }
+    result.sort(String.CASE_INSENSITIVE_ORDER);
+    return result.toArray(new String[0]);
+  }
+}