From 347bd66e810ccaaaccbd0c480b1d5dd2456f42e6 Mon Sep 17 00:00:00 2001 From: Tagir Valeev Date: Thu, 12 Jan 2017 11:18:29 +0700 Subject: [PATCH] StreamApiMigrationInspection: enable flatMap with primitive type change --- .../StreamApiMigrationInspection.java | 17 +++++++++++------ .../streamMigration/TerminalBlock.java | 4 ---- .../afterFlatMapChangePrimitiveType.java | 13 +++++++++++++ .../beforeFlatMapChangePrimitiveType.java | 18 ++++++++++++++++++ 4 files changed, 42 insertions(+), 10 deletions(-) create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFlatMapChangePrimitiveType.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFlatMapChangePrimitiveType.java 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 a68f29c7f5d6..e53c497ee4e5 100644 --- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java +++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java @@ -837,17 +837,22 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo @Override public String createReplacement() { String operation = "flatMap"; - PsiType type = mySource.getVariable().getType(); - if(type instanceof PsiPrimitiveType && !type.equals(myVariable.getType())) { - if(type.equals(PsiType.INT)) { + PsiType inType = myVariable.getType(); + PsiType outType = mySource.getVariable().getType(); + String lambda = myVariable.getName() + " -> " + getStreamExpression(); + if(outType instanceof PsiPrimitiveType && !outType.equals(inType)) { + if(outType.equals(PsiType.INT)) { operation = "flatMapToInt"; - } else if(type.equals(PsiType.LONG)) { + } else if(outType.equals(PsiType.LONG)) { operation = "flatMapToLong"; - } else if(type.equals(PsiType.DOUBLE)) { + } else if(outType.equals(PsiType.DOUBLE)) { operation = "flatMapToDouble"; } } - return "." + operation + "(" + myVariable.getName() + " -> " + getStreamExpression() + ")"; + if(inType instanceof PsiPrimitiveType && !outType.equals(inType)) { + return ".mapToObj(" + lambda + ")." + operation + "(" + CommonClassNames.JAVA_UTIL_FUNCTION_FUNCTION + ".identity())"; + } + return "." + operation + "(" + lambda + ")"; } @NotNull 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 6e23e89f17e6..dbe651206ddb 100644 --- a/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java +++ b/java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java @@ -163,10 +163,6 @@ class TerminalBlock { StreamSource source = StreamSource.tryCreate(loopStatement); final PsiStatement body = loopStatement.getBody(); if(source == null || body == null) return null; - // flatMap from primitive to primitive is supported only if primitive types match - // otherwise it would be necessary to create bogus step like - // .mapToObj(var -> collection.stream()).flatMap(Function.identity()) - if(myVariable.getType() instanceof PsiPrimitiveType && !myVariable.getType().equals(source.getVariable().getType())) return null; FlatMapOp op = new FlatMapOp(source, myVariable); TerminalBlock withFlatMap = new TerminalBlock(this, op, source.getVariable(), body); if(!VariableAccessUtils.variableIsUsed(myVariable, body)) { diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFlatMapChangePrimitiveType.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFlatMapChangePrimitiveType.java new file mode 100644 index 000000000000..625d45fb19dc --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFlatMapChangePrimitiveType.java @@ -0,0 +1,13 @@ +// "Replace with count()" "true" +import java.util.Arrays; +import java.util.function.Function; +import java.util.stream.IntStream; +import java.util.stream.LongStream; +import java.util.stream.Stream; + +public class Test { + public void test() { + long count = IntStream.range(0, 10).mapToObj(i -> LongStream.range(0, i)).flatMapToLong(Function.identity()).mapToObj(l -> Stream.of("x", "y", "z")).flatMap(Function.identity()).flatMapToInt(s -> IntStream.range(0, s.length())).count(); + System.out.println(count); + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFlatMapChangePrimitiveType.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFlatMapChangePrimitiveType.java new file mode 100644 index 000000000000..98ac2240159b --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFlatMapChangePrimitiveType.java @@ -0,0 +1,18 @@ +// "Replace with count()" "true" +import java.util.Arrays; + +public class Test { + public void test() { + long count = 0; + for(int i=0; i<10; i++) { + for(long l = 0; l < i; l++) { + for(String s : Arrays.asList("x", "y", "z")) { + for(int k = 0; k < s.length(); k++) { + count++; + } + } + } + } + System.out.println(count); + } +}