diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithCollectFix.java b/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithCollectFix.java index 06d617b43670..4dd8e7c06070 100644 --- a/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithCollectFix.java +++ b/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithCollectFix.java @@ -24,6 +24,7 @@ import com.intellij.psi.*; import com.intellij.psi.codeStyle.JavaCodeStyleManager; import com.intellij.psi.codeStyle.SuggestedNameInfo; import com.intellij.psi.codeStyle.VariableKind; +import com.intellij.psi.util.PsiTreeUtil; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -120,6 +121,12 @@ class ReplaceWithCollectFix extends MigrateToStreamFix { } } } + PsiElement nextStatement = PsiTreeUtil.skipSiblingsForward(foreachStatement, PsiComment.class, PsiWhiteSpace.class); + String comparatorText = StreamApiMigrationInspection.tryExtractSortComparatorText(nextStatement, variable); + if(comparatorText != null) { + builder.append(".sorted(").append(comparatorText).append(")"); + nextStatement.delete(); + } String callText = builder.append(".collect(java.util.stream.Collectors.") .append(createInitializerReplacementText(qualifierExpression.getType(), initializer)) .append(")").toString(); diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java b/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java index f791192355b0..41466665e68b 100644 --- a/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java +++ b/java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java @@ -659,6 +659,45 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo } } + /** + * + * @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") + 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).resolve() != list) return null; + if(comparatorExpression == null || ExpressionUtils.isNullLiteral(comparatorExpression)) return ""; + return comparatorExpression.getText(); + } + @Nullable static PsiMethodCallExpression extractToArrayExpression(PsiForeachStatement statement, PsiMethodCallExpression expression) { // return collection.toArray() or collection.toArray(new Type[0]) or collection.toArray(new Type[collection.size()]); diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSorted.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSorted.java new file mode 100644 index 000000000000..aa7c38f3bab2 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSorted.java @@ -0,0 +1,15 @@ +// "Replace with collect" "true" +import java.util.*; +import java.util.stream.Collectors; + +public class Collect { + class Person { + String getName() { + return ""; + } + } + + void collectNames(List persons){ + List names = persons.stream().map(Person::getName).sorted(Comparator.comparing(Person::getName)).collect(Collectors.toList()); + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSortedNatural.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSortedNatural.java new file mode 100644 index 000000000000..6f6b72b8f506 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSortedNatural.java @@ -0,0 +1,10 @@ +// "Replace with collect" "true" +import java.util.*; +import java.util.stream.Collectors; + +public class Collect { + void collectNames(List persons){ + List names = persons.stream().map(String::toLowerCase).sorted().collect(Collectors.toList()); + // sort + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSorted.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSorted.java new file mode 100644 index 000000000000..1c616c54b23d --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSorted.java @@ -0,0 +1,18 @@ +// "Replace with collect" "true" +import java.util.*; + +public class Collect { + class Person { + String getName() { + return ""; + } + } + + void collectNames(List persons){ + List names = new ArrayList<>(); + for (Person person : persons) { + names.add(person.getName()); + } + Collections.sort(names, Comparator.comparing(Person::getName)); + } +} diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSortedNatural.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSortedNatural.java new file mode 100644 index 000000000000..ce66211e9ecb --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeCollectSortedNatural.java @@ -0,0 +1,13 @@ +// "Replace with collect" "true" +import java.util.*; + +public class Collect { + void collectNames(List persons){ + List names = new ArrayList<>(); + for (String person : persons) { + names.add(person.toLowerCase()); + } + // sort + Collections.sort(names); + } +}