java 1.8, stream api migration: collect (to be continued)

java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java

@@ -106,9 +106,14 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
                     }
                   });
   
-                  if (effectivelyFinal[0] && !isTrivial(body, statement.getIterationParameter(), iteratedValueType)) {
-                    holder.registerProblem(iteratedValue, "Can be replaced with foreach call",
-                                           ProblemHighlightType.GENERIC_ERROR_OR_WARNING, new ReplaceWithForeachCallFix());
+                  if (effectivelyFinal[0]) {
+                    if (isCollectCall(body)) {
+                      holder.registerProblem(iteratedValue, "Can be replaced with collect call",
+                                             ProblemHighlightType.GENERIC_ERROR_OR_WARNING, new ReplaceWithCollectCallFix());
+                    } else if (!isTrivial(body, statement.getIterationParameter(), iteratedValueType)) {
+                      holder.registerProblem(iteratedValue, "Can be replaced with foreach call",
+                                             ProblemHighlightType.GENERIC_ERROR_OR_WARNING, new ReplaceWithForeachCallFix());
+                    }
                   }
                 }
               }
@@ -121,10 +126,42 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     };
   }
 
+  private static boolean isCollectCall(PsiStatement body) {
+    final PsiMethodCallExpression methodCallExpression = extractAddCall(body);
+    if (methodCallExpression != null) {
+      final PsiReferenceExpression methodExpression = methodCallExpression.getMethodExpression();
+      final PsiExpression qualifierExpression = methodExpression.getQualifierExpression();
+      PsiClass qualifierClass = null;
+      if (qualifierExpression instanceof PsiReferenceExpression) {
+        qualifierClass = PsiUtil.resolveClassInType(qualifierExpression.getType());
+      } else if (qualifierExpression == null) {
+        final PsiClass enclosingClass = PsiTreeUtil.getParentOfType(body, PsiClass.class);
+        if (PsiUtil.getEnclosingStaticElement(body, enclosingClass) == null) {
+          qualifierClass = enclosingClass;
+        }
+      }
+
+      if (qualifierClass != null && 
+          InheritanceUtil.isInheritor(qualifierClass, false, CommonClassNames.JAVA_UTIL_COLLECTION)) {
+
+        final PsiElement resolve = methodExpression.resolve();
+        if (resolve instanceof PsiMethod &&
+            "add".equals(((PsiMethod)resolve).getName()) &&
+            ((PsiMethod)resolve).getParameterList().getParametersCount() == 1) {
+          final PsiExpression[] args = methodCallExpression.getArgumentList().getExpressions();
+          if (args.length == 1) {
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+  
   private static boolean isTrivial(PsiStatement body, PsiParameter parameter, PsiType iteratedValueType) {
     final PsiIfStatement ifStatement = extractIfStatement(body);
     //stream
-    if (ifStatement != null && ifStatement.getElseBranch() == null && ifStatement.getThenBranch() != null && 
+    if (ifStatement != null &&
         InheritanceUtil.isInheritor(iteratedValueType, CommonClassNames.JAVA_UTIL_COLLECTION)) {
       return false;
     }
@@ -157,11 +194,12 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
 
           String foreEachText = body.getText();
           String iterated = iteratedValue.getText();
-          if (ifStmt != null && ifStmt.getElseBranch() == null) {
+          if (ifStmt != null) {
             final PsiExpression condition = ifStmt.getCondition();
             if (condition != null) {
               final PsiStatement thenBranch = ifStmt.getThenBranch();
-              if (thenBranch != null && InheritanceUtil.isInheritor(iteratedValue.getType(), CommonClassNames.JAVA_UTIL_COLLECTION)) {
+              LOG.assertTrue(thenBranch != null);
+              if (InheritanceUtil.isInheritor(iteratedValue.getType(), CommonClassNames.JAVA_UTIL_COLLECTION)) {
                 body = thenBranch;
                 foreEachText = thenBranch.getText();
                 iterated += ".stream().filter(" + parameter.getName() + " -> " + condition.getText() +")";
@@ -191,6 +229,76 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     }
   }
 
+  private static class ReplaceWithCollectCallFix implements LocalQuickFix {
+    @NotNull
+    @Override
+    public String getName() {
+      return getFamilyName();
+    }
+
+    @NotNull
+    @Override
+    public String getFamilyName() {
+      return "Replace with collect";
+    }
+
+    @Override
+    public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
+      final PsiForeachStatement foreachStatement = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PsiForeachStatement.class);
+      if (foreachStatement != null) {
+        PsiStatement body = foreachStatement.getBody();
+        final PsiExpression iteratedValue = foreachStatement.getIteratedValue();
+        if (body != null && iteratedValue != null) {
+          final PsiParameter parameter = foreachStatement.getIterationParameter();
+
+          final PsiIfStatement ifStatement = extractIfStatement(body);
+          final PsiMethodCallExpression methodCallExpression = extractAddCall(body);
+          String iteration = iteratedValue.getText() + ".stream()";
+          if (ifStatement != null) {
+            final PsiExpression condition = ifStatement.getCondition();
+            if (condition != null) {
+              iteration += ".filter(" + parameter.getName() + " -> " + condition.getText() +")";
+            }
+          }
+          iteration +=".map(" + parameter.getName() + " -> " + methodCallExpression.getArgumentList().getExpressions()[0].getText() + ").collect(java.util.stream.Collectors.";
+
+          String variableName = null;
+          PsiExpression initializer = null;
+          final PsiExpression qualifierExpression = methodCallExpression.getMethodExpression().getQualifierExpression();
+          if (qualifierExpression instanceof PsiReferenceExpression) {
+            final PsiElement resolve = ((PsiReferenceExpression)qualifierExpression).resolve();
+            if (resolve instanceof PsiVariable) {
+              if (resolve instanceof PsiLocalVariable && foreachStatement.equals(PsiTreeUtil.skipSiblingsForward(resolve.getParent(), PsiWhiteSpace.class))) {
+                initializer = ((PsiVariable)resolve).getInitializer();
+              }
+              variableName = ((PsiVariable)resolve).getName() + ".";
+            }
+          } else if (qualifierExpression == null) {
+            variableName = "";
+          }
+          
+          if (initializer != null) {
+            final PsiType initializerType = initializer.getType();
+            final PsiClassType rawType = initializerType instanceof PsiClassType ? ((PsiClassType)initializerType).rawType() : null;
+            if (rawType != null && rawType.equalsToText(CommonClassNames.JAVA_UTIL_ARRAY_LIST)) {
+              iteration += "toList()";
+            } else if (rawType != null && rawType.equalsToText(CommonClassNames.JAVA_UTIL_HASH_SET)) {
+              iteration += "toSet()";
+            } else {
+              iteration += "toCollection(() -> " + initializer.getText() +")";
+            }
+            iteration += ")";
+            initializer.replace(JavaPsiFacade.getElementFactory(project).createExpressionFromText(iteration, foreachStatement));
+            foreachStatement.delete();
+          } else if (variableName != null){
+            iteration += "toList())";
+            foreachStatement.replace(JavaPsiFacade.getElementFactory(project).createStatementFromText(variableName + "addAll(" + iteration +");", foreachStatement));
+          }
+        }
+      }
+    }
+  }
+
   public static PsiIfStatement extractIfStatement(PsiStatement body) {
     PsiIfStatement ifStmt = null;
     if (body instanceof PsiIfStatement) {
@@ -201,6 +309,34 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
         ifStmt = (PsiIfStatement)statements[0];
       }
     }
-    return ifStmt;
+    if (ifStmt != null && ifStmt.getElseBranch() == null && ifStmt.getThenBranch() != null) {
+      return ifStmt;
+    }
+    return null;
+  }
+  
+  private static PsiMethodCallExpression extractAddCall(PsiStatement body) {
+    final PsiIfStatement ifStatement = extractIfStatement(body);
+    if (ifStatement != null) {
+      return extractAddCall(ifStatement.getThenBranch());
+    }
+    PsiExpressionStatement stmt = null;
+    if (body instanceof PsiBlockStatement) {
+      final PsiStatement[] statements = ((PsiBlockStatement)body).getCodeBlock().getStatements();
+      if (statements.length == 1 && statements[0] instanceof PsiExpressionStatement) {
+        stmt = (PsiExpressionStatement)statements[0];
+      }
+    }
+    else if (body instanceof PsiExpressionStatement) {
+      stmt = (PsiExpressionStatement)body;
+    }
+
+    if (stmt != null) {
+      final PsiExpression expression = stmt.getExpression();
+      if (expression instanceof PsiMethodCallExpression) {
+        return (PsiMethodCallExpression)expression;
+      }
+    }
+    return null;
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectArrayList.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<Person> persons){
+    List<String> names = persons.stream().map(person -> person.getName()).collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectArrayListAndFilter.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<Person> persons){
+    List<String> names = persons.stream().filter(person -> person != null).map(person -> person.getName()).collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectHashSet.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<Person> persons){
+    Set<String> names = persons.stream().map(person -> person.getName()).collect(Collectors.toSet());
+  }
+}

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

@@ -0,0 +1,16 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  Set<String> names = new HashSet<>();
+  void collectNames(List<Person> persons){
+      names.addAll(persons.stream().map(person -> person.getName()).collect(Collectors.toList()));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectLinkedHashSet.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<Person> persons){
+    Set<String> names = persons.stream().map(person -> person.getName()).collect(Collectors.toCollection(() -> new LinkedHashSet<>()));
+  }
+}

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

@@ -0,0 +1,15 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public abstract class Collect implements Collection<String>{
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+      addAll(persons.stream().map(person -> person.getName()).collect(Collectors.toList()));
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterCollectSetParameter.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<Person> persons, Set<String> names){
+      names.addAll(persons.stream().map(person -> person.getName()).collect(Collectors.toList()));
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    List<String> names = new ArrayList<>();
+    for (Person person : pers<caret>ons) {
+      names.add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,20 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    List<String> names = new ArrayList<>();
+    for (Person person : pers<caret>ons) {
+      if (person != null) {
+        names.add(person.getName());
+      }
+    }
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    Set<String> names = new HashSet<>();
+    for (Person person : pers<caret>ons) {
+      names.add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  Set<String> names = new HashSet<>();
+  void collectNames(List<Person> persons){
+    for (Person person : pers<caret>ons) {
+      names.add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    Set<String> names = new LinkedHashSet<>();
+    for (Person person : pers<caret>ons) {
+      names.add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,17 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public abstract class Collect implements Collection<String>{
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons){
+    for (Person person : pers<caret>ons) {
+      add(person.getName());
+    }
+  }
+}

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

@@ -0,0 +1,17 @@
+// "Replace with collect" "true"
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Collect {
+  class Person {
+    String getName() {
+      return "";
+    }
+  }
+
+  void collectNames(List<Person> persons, Set<String> names){
+    for (Person person : pers<caret>ons) {
+      names.add(person.getName());
+    }
+  }
+}

platform/core-api/src/com/intellij/psi/CommonClassNames.java

@@ -61,6 +61,7 @@ public interface CommonClassNames {
   @NonNls String JAVA_UTIL_LIST = "java.util.List";
   @NonNls String JAVA_UTIL_ARRAY_LIST = "java.util.ArrayList";
   @NonNls String JAVA_UTIL_SET = "java.util.Set";
+  @NonNls String JAVA_UTIL_HASH_SET = "java.util.HashSet";
   @NonNls String JAVA_UTIL_PROPERTIES = "java.util.Properties";
   @NonNls String JAVA_UTIL_PROPERTY_RESOURCE_BUNDLE = "java.util.PropertyResourceBundle";
   @NonNls String JAVA_UTIL_DATE = "java.util.Date";