diff --git a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
index 23c0968d4e0f..79ae3bb9a5f1 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
@@ -169,7 +169,7 @@ abstract class FunctionHelper {
@NotNull
@Contract(pure = true)
- static FunctionHelper hashMapSupplier(PsiType type) {
+ static FunctionHelper newObjectSupplier(PsiType type, String instanceClassName) {
return new FunctionHelper(type) {
PsiExpression myExpression;
@@ -181,7 +181,7 @@ abstract class FunctionHelper {
@Override
void transform(StreamToLoopReplacementContext context, String... argumentValues) {
LOG.assertTrue(argumentValues.length == 0);
- myExpression = context.createExpression("new java.util.HashMap<>()");
+ myExpression = context.createExpression("new "+instanceClassName+"<>()");
}
};
}
diff --git a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/TerminalOperation.java b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/TerminalOperation.java
index 1b14a67c3d0f..07989090d82c 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/TerminalOperation.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/TerminalOperation.java
@@ -19,8 +19,10 @@ import com.intellij.codeInspection.streamToLoop.StreamToLoopInspection.StreamToL
import com.intellij.codeInspection.util.OptionalUtil;
import com.intellij.psi.*;
import com.intellij.psi.util.InheritanceUtil;
+import com.intellij.psi.util.PsiTypesUtil;
import com.intellij.psi.util.PsiUtil;
import com.intellij.psi.util.TypeConversionUtil;
+import com.intellij.util.ArrayUtil;
import com.siyeh.ig.psiutils.BoolUtils;
import com.siyeh.ig.psiutils.ExpressionUtils;
import one.util.streamex.StreamEx;
@@ -32,6 +34,7 @@ import java.util.Arrays;
import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Function;
+import java.util.stream.Stream;
/**
* @author Tagir Valeev
@@ -162,14 +165,15 @@ abstract class TerminalOperation extends Operation {
switch (collectorName) {
case "toList":
if (collectorArgs.length != 0) return null;
- return AccumulatedTerminalOperation.toList(resultType);
+ return ToCollectionTerminalOperation.toList(resultType);
case "toSet":
if (collectorArgs.length != 0) return null;
- return AccumulatedTerminalOperation.toCollection(resultType, CommonClassNames.JAVA_UTIL_HASH_SET, "set");
+ return new ToCollectionTerminalOperation(resultType,
+ FunctionHelper.newObjectSupplier(resultType, CommonClassNames.JAVA_UTIL_HASH_SET), "set");
case "toCollection":
if (collectorArgs.length != 1) return null;
fn = FunctionHelper.create(collectorArgs[0], 0);
- return fn == null ? null : new ToCollectionTerminalOperation(fn);
+ return fn == null ? null : new ToCollectionTerminalOperation(resultType, fn, null);
case "toMap": {
if (collectorArgs.length < 2 || collectorArgs.length > 4) return null;
FunctionHelper key = FunctionHelper.create(collectorArgs[0], 1);
@@ -178,7 +182,7 @@ abstract class TerminalOperation extends Operation {
PsiExpression merger = collectorArgs.length > 2 ? collectorArgs[2] : null;
FunctionHelper supplier = collectorArgs.length == 4
? FunctionHelper.create(collectorArgs[3], 0)
- : FunctionHelper.hashMapSupplier(resultType);
+ : FunctionHelper.newObjectSupplier(resultType, CommonClassNames.JAVA_UTIL_HASH_MAP);
if(supplier == null) return null;
return new ToMapTerminalOperation(key, value, merger, supplier, resultType);
}
@@ -241,7 +245,7 @@ abstract class TerminalOperation extends Operation {
if (resultSubType == null) return null;
CollectorOperation downstreamCollector;
if (collectorArgs.length == 1) {
- downstreamCollector = AccumulatedTerminalOperation.toList(resultSubType).asCollector();
+ downstreamCollector = ToCollectionTerminalOperation.toList(resultSubType).asCollector();
}
else {
PsiExpression downstream = collectorArgs[collectorArgs.length - 1];
@@ -255,7 +259,7 @@ abstract class TerminalOperation extends Operation {
}
FunctionHelper supplier = collectorArgs.length == 3
? FunctionHelper.create(collectorArgs[1], 0)
- : FunctionHelper.hashMapSupplier(resultType);
+ : FunctionHelper.newObjectSupplier(resultType, CommonClassNames.JAVA_UTIL_HASH_MAP);
return new GroupByTerminalOperation(fn, supplier, resultType, downstreamCollector);
}
case "minBy":
@@ -281,6 +285,67 @@ abstract class TerminalOperation extends Operation {
return null;
}
+ /**
+ * Eliminates <? extends> wildcards which correspond to the first supplied superclass.
+ * If there are more than one superclass supplied, performs the same operation for the last generic argument.
+ * Do not touch unrelated generic arguments which don't map to superclass type parameters.
+ *
+ * E.g.:
+ * {@code
+ * (List, Collection) -> List
+ * (MyList, Collection) -> MyList (assuming MyList extends List)
+ * (HashMap, Map, Collection) -> HashMap>
+ * (HashMap, Map) -> HashMap>
+ * }
+ *
+ * @param type
+ * @param superClasses
+ * @return
+ */
+ @NotNull
+ static PsiType eliminateCollectionWildcards(PsiType type, String... superClasses) {
+ if(superClasses.length == 0) return type;
+ String superClass = superClasses[0];
+ PsiClass aClass = PsiTypesUtil.getPsiClass(type);
+ if(aClass == null) return type;
+ PsiTypeParameter[] parameters = aClass.getTypeParameters();
+
+ if(parameters.length == 0) return type;
+ PsiSubstitutor substitutor = ((PsiClassType)type).resolveGenerics().getSubstitutor();
+ PsiElementFactory factory = JavaPsiFacade.getElementFactory(aClass.getProject());
+ PsiClassType classType = factory.createType(aClass, Stream.of(parameters).map(factory::createType).toArray(PsiType[]::new));
+ PsiClass baseClass = JavaPsiFacade.getInstance(aClass.getProject()).findClass(superClass, aClass.getResolveScope());
+ if(baseClass == null) return type;
+ PsiTypeParameter[] baseClassTypeParameters = baseClass.getTypeParameters();
+
+ for (int idx = 0; idx < baseClassTypeParameters.length; idx++) {
+ PsiClass parameter = PsiTypesUtil.getPsiClass(PsiUtil.substituteTypeParameter(classType, superClass, idx, false));
+ if(parameter instanceof PsiTypeParameter) {
+ PsiType origType = PsiUtil.substituteTypeParameter(type, superClass, idx, false);
+ PsiType replacedType = origType;
+ if(origType instanceof PsiCapturedWildcardType) {
+ replacedType = ((PsiCapturedWildcardType)origType).getUpperBound();
+ }
+ if(origType instanceof PsiWildcardType) {
+ replacedType = ((PsiWildcardType)origType).getExtendsBound();
+ }
+ if (idx == baseClassTypeParameters.length - 1) {
+ replacedType = eliminateCollectionWildcards(replacedType, Arrays.copyOfRange(superClasses, 1, superClasses.length));
+ }
+ if(replacedType != origType) {
+ substitutor = substitutor.put((PsiTypeParameter)parameter, replacedType);
+ }
+ }
+ }
+ return factory.createType(aClass, substitutor);
+ }
+
+ @NotNull
+ static String eliminateCollectionWildcards(StreamToLoopReplacementContext context,
+ String typeText, String... superClasses) {
+ return eliminateCollectionWildcards(context.createType(typeText), superClasses).getCanonicalText();
+ }
+
static class ReduceTerminalOperation extends TerminalOperation {
private PsiExpression myIdentity;
private String myType;
@@ -511,6 +576,11 @@ abstract class TerminalOperation extends Operation {
default void registerUsedNames(Consumer usedNameConsumer) {}
String getSupplier();
String getAccumulator(String acc, String item);
+
+ // Returns array of base classes for which "? extends" should be eliminated from generic arguments
+ // Several elements are used for downstream collection, in this case the first element is java.util.Map
+ // and second one corresponds to the map value
+ default String[] getBaseClassChain() {return ArrayUtil.EMPTY_STRING_ARRAY;}
}
abstract static class CollectorBasedTerminalOperation extends TerminalOperation implements CollectorOperation {
@@ -528,7 +598,8 @@ abstract class TerminalOperation extends Operation {
@Override
String generate(StreamVariable inVar, StreamToLoopReplacementContext context) {
transform(context, inVar.getName());
- String acc = context.declareResult(myAccNameSupplier.apply(context), myType, getSupplier(), true);
+ String acc = context.declareResult(myAccNameSupplier.apply(context),
+ eliminateCollectionWildcards(context, myType, getBaseClassChain()), getSupplier(), true);
return getAccumulator(acc, inVar.getName());
}
@@ -603,17 +674,6 @@ abstract class TerminalOperation extends Operation {
return myUpdateTemplate.replace("{acc}", acc).replace("{item}", item);
}
- @NotNull
- static AccumulatedTerminalOperation toCollection(PsiType collectionType, String implementationType, String varName) {
- return new AccumulatedTerminalOperation(varName, collectionType.getCanonicalText(), "new " + implementationType + "<>()",
- "{acc}.add({item});");
- }
-
- @NotNull
- private static AccumulatedTerminalOperation toList(@NotNull PsiType resultType) {
- return toCollection(resultType, CommonClassNames.JAVA_UTIL_ARRAY_LIST, "list");
- }
-
@NotNull
static AccumulatedTerminalOperation summing(PsiType type) {
return new AccumulatedTerminalOperation("sum", type.getCanonicalText(), "0", "{acc}+={item};");
@@ -627,14 +687,24 @@ abstract class TerminalOperation extends Operation {
}
static class ToCollectionTerminalOperation extends CollectorBasedTerminalOperation {
- public ToCollectionTerminalOperation(FunctionHelper fn) {
- super(fn.getResultType(), context -> fn.suggestFinalOutputNames(context, null, "collection").get(0), fn);
+ public ToCollectionTerminalOperation(PsiType resultType, FunctionHelper fn, String desiredName) {
+ super(resultType.getCanonicalText(), context -> fn.suggestFinalOutputNames(context, desiredName, "collection").get(0), fn);
}
@Override
public String getAccumulator(String acc, String item) {
return acc+".add("+item+");\n";
}
+
+ @Override
+ public String[] getBaseClassChain() {
+ return new String[] {CommonClassNames.JAVA_UTIL_COLLECTION};
+ }
+
+ @NotNull
+ private static ToCollectionTerminalOperation toList(@NotNull PsiType resultType) {
+ return new ToCollectionTerminalOperation(resultType, FunctionHelper.newObjectSupplier(resultType, CommonClassNames.JAVA_UTIL_ARRAY_LIST), "list");
+ }
}
static class MinMaxTerminalOperation extends TerminalOperation {
@@ -715,6 +785,11 @@ abstract class TerminalOperation extends Operation {
myMerger = merger;
}
+ @Override
+ public String[] getBaseClassChain() {
+ return new String[] {CommonClassNames.JAVA_UTIL_MAP};
+ }
+
@Override
public void registerUsedNames(Consumer usedNameConsumer) {
super.registerUsedNames(usedNameConsumer);
@@ -778,6 +853,11 @@ abstract class TerminalOperation extends Operation {
myCollector = collector;
}
+ @Override
+ public String[] getBaseClassChain() {
+ return ArrayUtil.prepend(CommonClassNames.JAVA_UTIL_MAP, myCollector.getBaseClassChain());
+ }
+
@Override
public void registerUsedNames(Consumer usedNameConsumer) {
super.registerUsedNames(usedNameConsumer);
@@ -831,7 +911,9 @@ abstract class TerminalOperation extends Operation {
@Override
String generate(StreamVariable inVar, StreamToLoopReplacementContext context) {
- String map = context.declareResult("map", myResultType, "new java.util.HashMap<>()", true);
+ String resultType = eliminateCollectionWildcards(context, myResultType,
+ ArrayUtil.prepend(CommonClassNames.JAVA_UTIL_MAP, myCollector.getBaseClassChain()));
+ String map = context.declareResult("map", resultType, "new java.util.HashMap<>()", true);
myPredicate.transform(context, inVar.getName());
myCollector.transform(context, inVar.getName());
context.addInitStep(map+".put(false, "+myCollector.getSupplier()+");");
@@ -862,6 +944,11 @@ abstract class TerminalOperation extends Operation {
myMapper.suggestVariableName(inVar, 0);
}
+ @Override
+ public String[] getBaseClassChain() {
+ return myDownstreamCollector.getBaseClassChain();
+ }
+
@Override
CollectorOperation asCollector() {
return myDownstreamCollector == null ? null : this;
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtends2.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtends2.java
new file mode 100644
index 000000000000..803ae38abd2a
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtends2.java
@@ -0,0 +1,26 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class Main {
+ private List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ List> list = new ArrayList<>();
+ for (CharSequence charSequence : getList()) {
+ List charSequences = asList(charSequence);
+ list.add(charSequences);
+ }
+ List> res2 = list;
+ System.out.println(res2);
+ }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsGroupingBy.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsGroupingBy.java
new file mode 100644
index 000000000000..e5da4ceaf69d
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsGroupingBy.java
@@ -0,0 +1,25 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+ static class MyList extends ArrayList {}
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ public MyList createList() {
+ return new MyList<>();
+ }
+
+ private void collect() {
+ Map> result = new HashMap<>();
+ for (CharSequence x : getList()) {
+ result.computeIfAbsent(x.length(), k -> createList()).add(x);
+ }
+ Map> map =
+ result;
+ System.out.println(map);
+ }
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsMyList.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsMyList.java
new file mode 100644
index 000000000000..97e60a2dff2f
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsMyList.java
@@ -0,0 +1,29 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class Main {
+ static class MyList extends ArrayList> {
+ }
+
+ private List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ MyList> res2 =
+ new MyList<>();
+ for (CharSequence charSequence : getList()) {
+ List charSequences = asList(charSequence);
+ res2.add(charSequences);
+ }
+ System.out.println(res2);
+ }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsToMap.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsToMap.java
new file mode 100644
index 000000000000..97b9629f1bbb
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterCollectExtendsToMap.java
@@ -0,0 +1,28 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class Main {
+ public List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ Map> result = new HashMap<>();
+ for (CharSequence charSequence : getList()) {
+ if (Objects.nonNull(charSequence)) {
+ if (result.put(charSequence, asList(charSequence)) != null) {
+ throw new IllegalStateException("Duplicate key");
+ }
+ }
+ }
+ Map> map = result;
+ System.out.println(map);
+ }
+}
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtends2.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtends2.java
new file mode 100644
index 000000000000..2c8c5c60e466
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtends2.java
@@ -0,0 +1,21 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class Main {
+ private List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ List> res2 = getList().stream().map(this::asList).collect(Collectors.toList());
+ System.out.println(res2);
+ }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsGroupingBy.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsGroupingBy.java
new file mode 100644
index 000000000000..dc18902fc107
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsGroupingBy.java
@@ -0,0 +1,24 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class Main {
+ static class MyList extends ArrayList {}
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ public MyList createList() {
+ return new MyList<>();
+ }
+
+ private void collect() {
+ Map> map =
+ getList().stream().collect(Collectors.groupingBy(x -> x.length(), Collectors.toCollection(this::createList)));
+ System.out.println(map);
+ }
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsMyList.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsMyList.java
new file mode 100644
index 000000000000..2d4b8c592b9c
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsMyList.java
@@ -0,0 +1,25 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class Main {
+ static class MyList extends ArrayList> {
+ }
+
+ private List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ MyList> res2 =
+ getList().stream().map(this::asList).collect(Collectors.toCollection(MyList::new));
+ System.out.println(res2);
+ }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsToMap.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsToMap.java
new file mode 100644
index 000000000000..5bccc535cd80
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeCollectExtendsToMap.java
@@ -0,0 +1,21 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class Main {
+ public List asList(CharSequence s) {
+ return Collections.singletonList(s);
+ }
+
+ public List getList() {
+ return Collections.emptyList();
+ }
+
+ private void collect() {
+ Map> map = getList()
+ .stream().filter(Objects::nonNull).collect(Collectors.toMap(Function.identity(), this::asList));
+ System.out.println(map);
+ }
+}