Convert for-loop with Collections.addAll: support primitive types; switch off when count expression is used; wrong contracts removed

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

@@ -301,13 +301,17 @@ class CollectMigration extends BaseStreamApiMigration {
       }
       String method = MethodCallUtils.isVarArgCall(myAddAllCall) ? CommonClassNames.JAVA_UTIL_STREAM_STREAM + "." + generic + "of"
                                                                  : CommonClassNames.JAVA_UTIL_ARRAYS + "." + generic + "stream";
-      return ".flatMap(" + myElement.getName() + "->" + method + "(" +
-             StreamEx.of(myAddAllCall.getArgumentList().getExpressions()).skip(1).map(PsiExpression::getText).joining(",") + "))";
+      String lambda = myElement.getName() + "->" + method + "(" +
+                      StreamEx.of(myAddAllCall.getArgumentList().getExpressions()).skip(1).map(PsiExpression::getText).joining(",") + ")";
+      return myElement.getType() instanceof PsiPrimitiveType ?
+             ".mapToObj(" + lambda + ").flatMap("+ CommonClassNames.JAVA_UTIL_FUNCTION_FUNCTION+".identity())" :
+             ".flatMap(" + lambda + ")";
     }
 
     @Nullable
     static AddingAllTerminal tryExtractAddAll(TerminalBlock tb, PsiMethodCallExpression call) {
-      if(!MethodCallUtils.isCallToStaticMethod(call, CommonClassNames.JAVA_UTIL_COLLECTIONS, "addAll", 2)) {
+      if(tb.getCountExpression() != null ||
+         !MethodCallUtils.isCallToStaticMethod(call, CommonClassNames.JAVA_UTIL_COLLECTIONS, "addAll", 2)) {
         return null;
       }
       PsiExpression[] args = call.getArgumentList().getExpressions();
@@ -501,7 +505,7 @@ class CollectMigration extends BaseStreamApiMigration {
       myStatement.delete();
     }
 
-    @Contract("null, _ -> null")
+    @Nullable
     public static CollectTerminal tryWrap(CollectTerminal terminal, PsiElement element) {
       PsiVariable list = terminal.getTargetVariable();
       if (list == null || !(element instanceof PsiExpressionStatement)) return null;
@@ -606,7 +610,7 @@ class CollectMigration extends BaseStreamApiMigration {
       myUpstream.cleanUp();
     }
 
-    @Contract("null, _, _ -> null")
+    @Nullable
     public static ToArrayTerminal tryWrap(CollectTerminal terminal, PsiLoopStatement loopStatement, PsiElement element) {
       PsiVariable collectionVariable = terminal.getTargetVariable();
       if (collectionVariable == null || StreamApiMigrationInspection.getInitializerUsageStatus(collectionVariable, loopStatement)

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

@@ -0,0 +1,11 @@
+// "Replace with toArray" "true"
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+public class Test {
+  String[] test(int count) {
+      return IntStream.range(0, count).mapToObj(i -> Stream.of("one", "two", "three")).flatMap(Function.identity()).toArray(String[]::new);
+  }
+}

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

@@ -0,0 +1,12 @@
+// "Replace with toArray" "true"
+import java.util.*;
+
+public class Test {
+  String[] test(int count) {
+    List<String> result = new ArrayList<>();
+    for(int <caret>i=0; i < count; i++) {
+      Collections.addAll(result, "one", "two", "three");
+    }
+    return result.toArray(new String[result.size()]);
+  }
+}

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

@@ -0,0 +1,16 @@
+// "Replace with toArray" "false"
+import java.util.*;
+
+public class Test {
+  Object[] test(List<String[]> list) {
+    List<Object> result = new LinkedList<>();
+    for(String[] str : li<caret>st) {
+      if(str != null) {
+        Collections.addAll(result, str);
+        if(result.size() > 10) break;
+      }
+    }
+    result.sort(null);
+    return result.toArray();
+  }
+}