IDEA-165397 Replace stream API with loop: support lambdas as toArray argument

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

@@ -181,31 +181,31 @@ abstract class FunctionHelper {
   }
 
   /**
-   * Renames references to the variable oldName in given expression into newName
+   * Replaces all the references to the variable {@code name} in given expression with {@code replacement}.
+   *
+   * <p>
+   *   If the replacement is a new name to the variable, the caller must take care that this new name was not used before.
+   * </p>
+   *
    * @param expression an expression to search-and-replace references inside
-   * @param oldName old name
-   * @param newName new name
+   * @param name a reference name to replace
+   * @param replacement a replacement expression (new name or literal)
    * @param context context
-   * @return resulting expression (might be the same as input expression) or null if expression already had references to newName,
-   *   so rename may merge two variables
+   * @return resulting expression (might be the same as input expression)
    */
   @NotNull
-  static PsiExpression renameVarReference(@NotNull PsiExpression expression,
-                                          String oldName,
-                                          String newName,
-                                          StreamToLoopReplacementContext context) {
-    if(oldName.equals(newName)) return expression;
-    PsiLambdaExpression lambda = (PsiLambdaExpression)context.createExpression("("+oldName+","+newName+")-> "+expression.getText());
-    PsiParameter[] parameters = lambda.getParameterList().getParameters();
-    PsiParameter oldVar = parameters[0];
-    PsiParameter newVar = parameters[1];
+  static PsiExpression replaceVarReference(@NotNull PsiExpression expression,
+                                           String name,
+                                           String replacement,
+                                           StreamToLoopReplacementContext context) {
+    if(name.equals(replacement)) return expression;
+    PsiLambdaExpression lambda = (PsiLambdaExpression)context.createExpression(name+"->"+expression.getText());
+    PsiParameter var = lambda.getParameterList().getParameters()[0];
     PsiElement body = lambda.getBody();
     LOG.assertTrue(body != null);
-    if(ReferencesSearch.search(newVar, new LocalSearchScope(body)).findFirst() != null) {
-      throw new IllegalStateException("Reference with name "+newVar+" already exists in "+lambda.getText());
-    }
-    for (PsiReference ref : ReferencesSearch.search(oldVar, new LocalSearchScope(body)).findAll()) {
-      ref.handleElementRename(newName);
+    PsiExpression replacementExpression = context.createExpression(replacement);
+    for (PsiReference ref : ReferencesSearch.search(var, new LocalSearchScope(body)).findAll()) {
+      ref.getElement().replace(replacementExpression);
     }
     return (PsiExpression)lambda.getBody();
   }
@@ -304,7 +304,7 @@ abstract class FunctionHelper {
       LOG.assertTrue(body instanceof PsiExpression);
       myExpression = (PsiExpression)body;
       EntryStream.zip(lambda.getParameterList().getParameters(), newNames)
-        .forKeyValue((param, newName) -> myExpression = renameVarReference(myExpression, param.getName(), newName, context));
+        .forKeyValue((param, newName) -> myExpression = replaceVarReference(myExpression, param.getName(), newName, context));
     }
 
     @Override
@@ -336,7 +336,7 @@ abstract class FunctionHelper {
       if(oldName.equals(newName)) return;
       PsiExpression qualifier = myMethodRef.getQualifierExpression();
       if(qualifier == null) return;
-      qualifier = renameVarReference(qualifier, oldName, newName, context);
+      qualifier = replaceVarReference(qualifier, oldName, newName, context);
       myMethodRef = fromText(context, qualifier.getText()+"::"+myMethodRef.getReferenceName());
     }
   }
@@ -382,7 +382,7 @@ abstract class FunctionHelper {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myExpression = renameVarReference(myExpression, oldName, newName, context);
+      myExpression = replaceVarReference(myExpression, oldName, newName, context);
     }
 
     @Override
@@ -421,7 +421,7 @@ abstract class FunctionHelper {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myReference = renameVarReference(myReference, oldName, newName, context);
+      myReference = replaceVarReference(myReference, oldName, newName, context);
     }
 
     @Override
@@ -473,7 +473,7 @@ abstract class FunctionHelper {
     void transform(StreamToLoopReplacementContext context, String... newNames) {
       LOG.assertTrue(newNames.length == myParameters.length);
       EntryStream.zip(myParameters, newNames).forKeyValue(
-        (oldName, newName) -> myBody = renameVarReference(myBody, oldName, newName, context));
+        (oldName, newName) -> myBody = replaceVarReference(myBody, oldName, newName, context));
     }
 
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
@@ -484,7 +484,7 @@ abstract class FunctionHelper {
           if (!paramName.equals(oldName) &&
               !StreamEx.of(myParameters).has(paramName)) {
             try {
-              myBody = renameVarReference(myBody, newName, paramName, context);
+              myBody = replaceVarReference(myBody, newName, paramName, context);
               myParameters[(int)idx.getAsLong()] = paramName;
               break;
             }
@@ -496,7 +496,7 @@ abstract class FunctionHelper {
           }
         }
       }
-      myBody = renameVarReference(myBody, oldName, newName, context);
+      myBody = replaceVarReference(myBody, oldName, newName, context);
     }
 
     @Override

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

@@ -264,7 +264,7 @@ abstract class Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myExpression = FunctionHelper.renameVarReference(myExpression, oldName, newName, context);
+      myExpression = FunctionHelper.replaceVarReference(myExpression, oldName, newName, context);
     }
 
     @Override
@@ -288,7 +288,7 @@ abstract class Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myLimit = FunctionHelper.renameVarReference(myLimit, oldName, newName, context);
+      myLimit = FunctionHelper.replaceVarReference(myLimit, oldName, newName, context);
     }
 
     @Override

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

@@ -30,7 +30,7 @@ import java.util.Arrays;
 import java.util.function.Consumer;
 
 import static com.intellij.codeInspection.streamToLoop.FunctionHelper.processUsedNames;
-import static com.intellij.codeInspection.streamToLoop.FunctionHelper.renameVarReference;
+import static com.intellij.codeInspection.streamToLoop.FunctionHelper.replaceVarReference;
 
 /**
  * @author Tagir Valeev
@@ -109,7 +109,7 @@ abstract class SourceOperation extends Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myQualifier = renameVarReference(myQualifier, oldName, newName, context);
+      myQualifier = replaceVarReference(myQualifier, oldName, newName, context);
     }
 
     @Override
@@ -146,7 +146,7 @@ abstract class SourceOperation extends Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      Arrays.asList(myArgList).replaceAll(arg -> renameVarReference(arg, oldName, newName, context));
+      Arrays.asList(myArgList).replaceAll(arg -> replaceVarReference(arg, oldName, newName, context));
     }
 
     @Override
@@ -189,7 +189,7 @@ abstract class SourceOperation extends Operation {
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
       myFn.rename(oldName, newName, context);
       if(myLimit != null) {
-        myLimit = renameVarReference(myLimit, oldName, newName, context);
+        myLimit = replaceVarReference(myLimit, oldName, newName, context);
       }
     }
 
@@ -227,7 +227,7 @@ abstract class SourceOperation extends Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myInitializer = renameVarReference(myInitializer, oldName, newName, context);
+      myInitializer = replaceVarReference(myInitializer, oldName, newName, context);
       myFn.rename(oldName, newName, context);
     }
 
@@ -264,8 +264,8 @@ abstract class SourceOperation extends Operation {
 
     @Override
     void rename(String oldName, String newName, StreamToLoopReplacementContext context) {
-      myOrigin = renameVarReference(myOrigin, oldName, newName, context);
-      myBound = renameVarReference(myBound, oldName, newName, context);
+      myOrigin = replaceVarReference(myOrigin, oldName, newName, context);
+      myBound = replaceVarReference(myBound, oldName, newName, context);
     }
 
     @Override

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

@@ -111,20 +111,12 @@ abstract class TerminalOperation extends Operation {
         if(args.length == 0) return new ToPrimitiveArrayTerminalOperation(componentType.getCanonicalText());
       }
       else {
-        String arr = "";
+        FunctionHelper fn = null;
         if(args.length == 1) {
-          if(!(args[0] instanceof PsiMethodReferenceExpression)) return null;
-          PsiMethodReferenceExpression arrCtor = (PsiMethodReferenceExpression)args[0];
-          if(!arrCtor.isConstructor()) return null;
-          PsiTypeElement typeElement = arrCtor.getQualifierType();
-          if(typeElement == null) return null;
-          PsiType type = typeElement.getType();
-          if(!(type instanceof PsiArrayType)) return null;
-          arr = "new "+type.getCanonicalText().replaceFirst("\\[]", "[0]");
+          fn = FunctionHelper.create(args[0], 1);
+          if(fn == null) return null;
         }
-        return new AccumulatedTerminalOperation("list", CommonClassNames.JAVA_UTIL_LIST + "<" + elementType.getCanonicalText() + ">",
-                                                "new "+ CommonClassNames.JAVA_UTIL_ARRAY_LIST+"<>()", "{acc}.add({item});",
-                                                "{acc}.toArray("+arr+")");
+        return new ToArrayTerminalOperation(elementType, fn);
       }
     }
     if ((name.equals("max") || name.equals("min")) && args.length < 2) {
@@ -425,6 +417,29 @@ abstract class TerminalOperation extends Operation {
     }
   }
 
+  static class ToArrayTerminalOperation extends TerminalOperation {
+    private final String myType;
+    private final FunctionHelper mySupplier;
+
+    public ToArrayTerminalOperation(PsiType type, FunctionHelper supplier) {
+      myType = type.getCanonicalText();
+      mySupplier = supplier;
+    }
+
+    @Override
+    String generate(StreamVariable inVar, StreamToLoopReplacementContext context) {
+      String list = context.declareResult("list", CommonClassNames.JAVA_UTIL_LIST + "<" + myType + ">",
+                                          "new " + CommonClassNames.JAVA_UTIL_ARRAY_LIST + "<>()", false);
+      String toArrayArg = "";
+      if(mySupplier != null) {
+        mySupplier.transform(context, "0");
+        toArrayArg = mySupplier.getText();
+      }
+      context.setFinisher(list + ".toArray(" + toArrayArg + ")");
+      return list+".add("+inVar+");\n";
+    }
+  }
+
   static class FindTerminalOperation extends TerminalOperation {
     private String myType;
 

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

@@ -0,0 +1,16 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.function.*;
+
+public class Main {
+  private static <A> A[] toArraySkippingNulls(List<?> list, IntFunction<A[]> generator) {
+      List<Object> result = new ArrayList<>();
+      for (Object o : list) {
+          if (Objects.nonNull(o)) {
+              result.add(o);
+          }
+      }
+      return result.toArray(generator.apply(0));
+  }
+}

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

@@ -0,0 +1,20 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public class Main {
+  private static Number[] test(int[] numbers) {
+      List<Integer> list = new ArrayList<>();
+      for (int number : numbers) {
+          Integer integer = number;
+          list.add(integer);
+      }
+      return list.toArray(new Integer[0]);
+  }
+
+  public static void main(String[] args) {
+    System.out.println(Arrays.asList(test(new int[] {1,2,3})));
+  }
+}

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

@@ -0,0 +1,10 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.function.*;
+
+public class Main {
+  private static <A> A[] toArraySkippingNulls(List<?> list, IntFunction<A[]> generator) {
+    return list.stream().filter(Objects::nonNull).toAr<caret>ray(generator);
+  }
+}

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

@@ -0,0 +1,13 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.Arrays;
+
+public class Main {
+  private static Number[] test(int[] numbers) {
+    return Arrays.stream(numbers).boxed().toArr<caret>ay(size -> new Integer[size]);
+  }
+
+  public static void main(String[] args) {
+    System.out.println(Arrays.asList(test(new int[] {1,2,3})));
+  }
+}