StreamToLoopInspection: sorted() operation (currently for non-primitive streams only)

java/java-impl/src/com/intellij/codeInspection/streamToLoop/Operation.java

@@ -16,15 +16,13 @@
 package com.intellij.codeInspection.streamToLoop;
 
 import com.intellij.codeInspection.streamToLoop.StreamToLoopInspection.StreamToLoopReplacementContext;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiExpression;
-import com.intellij.psi.PsiMethodCallExpression;
-import com.intellij.psi.PsiType;
+import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTypesUtil;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
 
@@ -71,6 +69,9 @@ abstract class Operation {
       FunctionHelper fn = FunctionHelper.create(args[0], 1);
       return fn == null ? null : new FilterOperation(fn);
     }
+    if(name.equals("sorted") && !(inType instanceof PsiPrimitiveType)) {
+      return new SortedOperation(args.length == 1 ? args[0] : null);
+    }
     if(name.equals("peek") && args.length == 1) {
       FunctionHelper fn = FunctionHelper.create(args[0], 1);
       return fn == null ? null : new PeekOperation(fn);
@@ -229,7 +230,7 @@ abstract class Operation {
       for(StreamToLoopInspection.OperationRecord or : StreamEx.ofReversed(myRecords)) {
         replacement = or.myOperation.wrap(or.myInVar, or.myOutVar, replacement, innerContext);
       }
-      return StreamEx.of(innerContext.getDeclarations()).map(str -> str  + "\n").joining()+replacement;
+      return replacement;
     }
 
     @Nullable
@@ -303,4 +304,23 @@ abstract class Operation {
       return "if(" + limit + "--==0) " + context.getBreakStatement() + code;
     }
   }
+
+  static class SortedOperation extends Operation {
+    private final @Nullable PsiExpression myComparator;
+
+    SortedOperation(@Nullable PsiExpression comparator) {
+      myComparator = comparator;
+    }
+
+    @Override
+    String wrap(StreamVariable inVar, StreamVariable outVar, String code, StreamToLoopReplacementContext context) {
+      String list = context.registerVarName(Arrays.asList("toSort", "listToSort"));
+      context.addAfterStep(new SourceOperation.ForEachSource(context.createExpression(list)).wrap(null, outVar, code, context));
+      context.addAfterStep(list + ".sort(" + (myComparator == null ? "null" : myComparator.getText()) + ");\n");
+      String listType = CommonClassNames.JAVA_UTIL_LIST + "<" + inVar.getType() + ">";
+      String initializer = "new " + CommonClassNames.JAVA_UTIL_ARRAY_LIST + "<>()";
+      context.addBeforeStep(listType + " " + list + "=" + initializer + ";");
+      return list+".add("+inVar+");\n";
+    }
+  }
 }

java/java-impl/src/com/intellij/codeInspection/streamToLoop/SourceOperation.java

@@ -24,6 +24,7 @@ import com.siyeh.ig.psiutils.ExpressionUtils;
 import com.siyeh.ig.psiutils.StreamApiUtil;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
@@ -41,9 +42,12 @@ abstract class SourceOperation extends Operation {
     return true;
   }
 
+  @NotNull
   @Override
   final String wrap(StreamVariable inVar, StreamVariable outVar, String code, StreamToLoopReplacementContext context) {
-    return wrap(outVar, code, context);
+    // Cannot inline "result" as wrap may register more beforeSteps
+    String result = wrap(outVar, code, context);
+    return context.drainBeforeSteps() + result + context.drainAfterSteps();
   }
 
   abstract String wrap(StreamVariable outVar, String code, StreamToLoopReplacementContext context);

java/java-impl/src/com/intellij/codeInspection/streamToLoop/StreamToLoopInspection.java

@@ -257,9 +257,6 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
           replacement = or.myOperation.wrap(or.myInVar, or.myOutVar, replacement, context);
         }
         ct.insertCommentsBefore(statement);
-        for (String declaration : context.getDeclarations()) {
-          addStatement(project, statement, factory.createStatementFromText(declaration, statement));
-        }
         for (PsiStatement addedStatement : ((PsiBlockStatement)factory.createStatementFromText("{" + replacement + "}", statement))
           .getCodeBlock().getStatements()) {
           addStatement(project, statement, addedStatement);
@@ -360,7 +357,8 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
     private final PsiStatement myStatement;
     private final Set<String> myUsedNames;
     private final Set<String> myUsedLabels;
-    private final List<String> myDeclarations = new ArrayList<>();
+    private final List<String> myBeforeSteps = new ArrayList<>();
+    private final List<String> myAfterSteps = new ArrayList<>();
     private final CommentTracker myCommentTracker;
     private PsiElement myPlaceholder;
     private final PsiElementFactory myFactory;
@@ -425,10 +423,6 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
       return label == null ? "break;\n" : "break "+label+";\n";
     }
 
-    public List<String> getDeclarations() {
-      return myDeclarations;
-    }
-
     public String registerVarName(Collection<String> variants) {
       if(variants.isEmpty()) {
         return registerVarName(Collections.singleton("val"));
@@ -452,12 +446,28 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
     public String declare(String desiredName, String type, String initializer) {
       String name = registerVarName(
         mySuffix.isEmpty() ? Collections.singleton(desiredName) : Arrays.asList(desiredName, desiredName + mySuffix));
-      myDeclarations.add(type + " " + name + " = " + initializer + ";");
+      myBeforeSteps.add(type + " " + name + " = " + initializer + ";");
       return name;
     }
 
-    public void addInitStep(String initStatement) {
-      myDeclarations.add(initStatement);
+    public void addBeforeStep(String beforeStatement) {
+      myBeforeSteps.add(beforeStatement);
+    }
+
+    public void addAfterStep(String afterStatement) {
+      myAfterSteps.add(0, afterStatement);
+    }
+
+    public String drainAfterSteps() {
+      String afterSteps = String.join("", myAfterSteps);
+      myAfterSteps.clear();
+      return afterSteps;
+    }
+
+    public String drainBeforeSteps() {
+      String beforeSteps = String.join("", myBeforeSteps);
+      myBeforeSteps.clear();
+      return beforeSteps;
     }
 
     public String declareResult(String desiredName, String type, String initializer, @NotNull ResultKind kind) {
@@ -472,13 +482,13 @@ public class StreamToLoopInspection extends BaseJavaBatchLocalInspectionTool {
             PsiExpression oldInitializer = copy.getInitializer();
             LOG.assertTrue(oldInitializer != null);
             oldInitializer.replace(createExpression(initializer));
-            myDeclarations.add(copy.getText());
+            myBeforeSteps.add(copy.getText());
             return var.getName();
           }
         }
       }
       String name = registerVarName(Arrays.asList(desiredName, "result"));
-      myDeclarations.add(type + " " + name + " = " + initializer + ";");
+      myBeforeSteps.add(type + " " + name + " = " + initializer + ";");
       if(myFinisher != null) {
         throw new IllegalStateException("Finisher is already defined");
       }

java/java-impl/src/com/intellij/codeInspection/streamToLoop/TerminalOperation.java

@@ -887,8 +887,8 @@ abstract class TerminalOperation extends Operation {
       String map = context.declareResult("map", resultType.getCanonicalText(), "new java.util.HashMap<>()", ResultKind.FINAL);
       myPredicate.transform(context, inVar.getName());
       myCollector.transform(context, inVar.getName());
-      context.addInitStep(map + ".put(false, " + myCollector.getSupplier() + ");");
-      context.addInitStep(map + ".put(true, " + myCollector.getSupplier() + ");");
+      context.addBeforeStep(map + ".put(false, " + myCollector.getSupplier() + ");");
+      context.addBeforeStep(map + ".put(true, " + myCollector.getSupplier() + ");");
       return myCollector.getAccumulator(map + ".get(" + myPredicate.getText() + ")", inVar.getName());
     }
   }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterFlatMapSorted.java

@@ -0,0 +1,23 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+
+public class Main {
+  public String testSorted(List<List<String>> list) {
+      for (List<String> lst : list) {
+          List<String> toSort = new ArrayList<>();
+          for (String x : lst) {
+              if (x != null) {
+                  toSort.add(x);
+              }
+          }
+          toSort.sort(null);
+          for (String x : toSort) {
+              if (x.length() < 5) {
+                  return x;
+              }
+          }
+      }
+      return "";
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterSorted.java

@@ -0,0 +1,25 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+  public List<String> testSorted(List<String> list) {
+      List<String> toSort = new ArrayList<>();
+      for (String s : list) {
+          if (s != null) {
+              toSort.add(s);
+          }
+      }
+      toSort.sort(null);
+      List<String> result = new ArrayList<>();
+      Set<String> uniqueValues = new HashSet<>();
+      for (String s : toSort) {
+          String trim = s.trim();
+          if (uniqueValues.add(trim)) {
+              result.add(trim);
+          }
+      }
+      return result;
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterSortedSimple.java

@@ -0,0 +1,19 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+  public List<String> testSorted(List<String> list) {
+      List<String> toSort = new ArrayList<>();
+      for (String s : list) {
+          toSort.add(s);
+      }
+      toSort.sort(String.CASE_INSENSITIVE_ORDER);
+      List<String> result = new ArrayList<>();
+      for (String s : toSort) {
+          result.add(s);
+      }
+      return result;
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeFlatMapSorted.java

@@ -0,0 +1,9 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+
+public class Main {
+  public String testSorted(List<List<String>> list) {
+    return list.stream().flatMap(lst -> lst.stream().filter(Objects::nonNull).sorted()).filter(x -> x.length() < 5).<caret>findFirst().orElse("");
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeSorted.java

@@ -0,0 +1,10 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+  public List<String> testSorted(List<String> list) {
+    return list.stream().filter(Objects::nonNull).sorted().map(String::trim).distinct().<caret>collect(Collectors.toList());
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeSortedSimple.java

@@ -0,0 +1,10 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.*;
+
+public class Main {
+  public List<String> testSorted(List<String> list) {
+    return list.stream().sorted(String.CASE_INSENSITIVE_ORDER).<caret>collect(Collectors.toList());
+  }
+}