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 68f5f2ff61da..b0caeb0ba084 100644 --- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/CollectMigration.java +++ b/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) diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectionsAddAllPrimitive.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectionsAddAllPrimitive.java new file mode 100644 index 000000000000..8d33c6a5ea9f --- /dev/null +++ b/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); + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectionsAddAllPrimitive.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectionsAddAllPrimitive.java new file mode 100644 index 000000000000..64f6ba3d3ee1 --- /dev/null +++ b/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 result = new ArrayList<>(); + for(int i=0; i < count; i++) { + Collections.addAll(result, "one", "two", "three"); + } + return result.toArray(new String[result.size()]); + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectionsAddAllWrongLimit.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectionsAddAllWrongLimit.java new file mode 100644 index 000000000000..84cff2afd8c8 --- /dev/null +++ b/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 list) { + List result = new LinkedList<>(); + for(String[] str : list) { + if(str != null) { + Collections.addAll(result, str); + if(result.size() > 10) break; + } + } + result.sort(null); + return result.toArray(); + } +}