IDEA-170626 Suggest to use Stream.peek instead of Stream.map if the lambda parameter is not reassigned during the operation.

java/java-impl/src/com/intellij/codeInspection/SimplifyStreamApiCallChainsInspection.java

@@ -61,6 +61,8 @@ public class SimplifyStreamApiCallChainsInspection extends BaseJavaBatchLocalIns
     instanceCall(CommonClassNames.JAVA_UTIL_STREAM_STREAM, "filter").parameterTypes(CommonClassNames.JAVA_UTIL_FUNCTION_PREDICATE);
   private static final CallMatcher STREAM_MAP =
     instanceCall(CommonClassNames.JAVA_UTIL_STREAM_STREAM, "map").parameterTypes(CommonClassNames.JAVA_UTIL_FUNCTION_FUNCTION);
+  private static final CallMatcher BASE_STREAM_MAP =
+    instanceCall(CommonClassNames.JAVA_UTIL_STREAM_BASE_STREAM, "map").parameterCount(1);
   private static final CallMatcher STREAM_ANY_MATCH =
     instanceCall(CommonClassNames.JAVA_UTIL_STREAM_BASE_STREAM, "anyMatch").parameterCount(1);
   private static final CallMatcher STREAM_NONE_MATCH =
@@ -81,7 +83,8 @@ public class SimplifyStreamApiCallChainsInspection extends BaseJavaBatchLocalIns
     ReplaceWithBoxedFix.handler(),
     ReplaceWithElementIterationFix.handler(),
     ReplaceForEachMethodFix.handler(),
-    RemoveBooleanIdentityFix.handler()
+    RemoveBooleanIdentityFix.handler(),
+    ReplaceWithPeekFix.handler()
   ).registerAll(SimplifyMatchNegationFix.handlers());
 
   private static final Logger LOG = Logger.getInstance("#" + SimplifyStreamApiCallChainsInspection.class.getName());
@@ -720,6 +723,53 @@ public class SimplifyStreamApiCallChainsInspection extends BaseJavaBatchLocalIns
     }
   }
 
+  private static class ReplaceWithPeekFix implements CallChainSimplification {
+
+    @Override
+    public String getName() {
+      return "Replace with 'peek'";
+    }
+
+    @Override
+    public String getMessage() {
+      return "Can be replaced with 'peek'";
+    }
+
+    @Override
+    public PsiElement simplify(PsiMethodCallExpression call) {
+      PsiLambdaExpression lambda =
+        tryCast(PsiUtil.skipParenthesizedExprDown(call.getArgumentList().getExpressions()[0]), PsiLambdaExpression.class);
+      if (lambda == null) return null;
+      PsiCodeBlock block = tryCast(lambda.getBody(), PsiCodeBlock.class);
+      if (block == null) return null;
+      PsiReturnStatement statement = tryCast(ArrayUtil.getLastElement(block.getStatements()), PsiReturnStatement.class);
+      if (statement == null) return null;
+      ExpressionUtils.bindCallTo(call, "peek");
+      new CommentTracker().deleteAndRestoreComments(statement);
+      LambdaRefactoringUtil.simplifyToExpressionLambda(lambda);
+      LambdaCanBeMethodReferenceInspection.replaceLambdaWithMethodReference(lambda);
+      return call;
+    }
+
+    static CallHandler<CallChainSimplification> handler() {
+      return CallHandler.of(BASE_STREAM_MAP, call -> {
+        PsiLambdaExpression lambda =
+          tryCast(PsiUtil.skipParenthesizedExprDown(call.getArgumentList().getExpressions()[0]), PsiLambdaExpression.class);
+        if (lambda == null) return null;
+        PsiParameter[] parameters = lambda.getParameterList().getParameters();
+        if (parameters.length != 1) return null;
+        PsiCodeBlock block = tryCast(lambda.getBody(), PsiCodeBlock.class);
+        if (block == null) return null;
+        PsiStatement[] statements = block.getStatements();
+        if (statements.length != 2) return null;
+        PsiReturnStatement returnStatement = tryCast(statements[1], PsiReturnStatement.class);
+        if (returnStatement == null || !ExpressionUtils.isReferenceTo(returnStatement.getReturnValue(), parameters[0])) return null;
+        if (VariableAccessUtils.variableIsAssigned(parameters[0]) || !(statements[0] instanceof PsiExpressionStatement)) return null;
+        return new ReplaceWithPeekFix();
+      });
+    }
+  }
+
   private static class ReplaceWithBoxedFix implements CallChainSimplification {
     private static final CallMatcher MAP_TO_OBJ = instanceCall(CommonClassNames.JAVA_UTIL_STREAM_BASE_STREAM, "mapToObj").parameterCount(1);
 

java/java-tests/testData/inspection/streamApiCallChains/afterMapToPeek.java

@@ -0,0 +1,14 @@
+// "Replace with 'peek'" "true"
+
+import java.util.List;
+
+public class Main {
+  void test(List<String> list) {
+      // hello
+/* in return */
+      long count = list.stream()
+      .peek(System.out::println)
+      .count();
+    System.out.println(count);
+  }
+}

java/java-tests/testData/inspection/streamApiCallChains/afterMapToPeekPrimitive.java

@@ -0,0 +1,15 @@
+// "Replace with 'peek'" "true"
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
+
+public class Main {
+  void test() {
+    AtomicInteger counter = new AtomicInteger();
+    int[] ints = IntStream.range(0, 100)
+      .filter(x -> x % 3 == 0)
+      .peek((x -> counter.incrementAndGet()))
+      .toArray();
+    System.out.println(counter.get());
+  }
+}

java/java-tests/testData/inspection/streamApiCallChains/beforeMapToPeek.java

@@ -0,0 +1,16 @@
+// "Replace with 'peek'" "true"
+
+import java.util.List;
+
+public class Main {
+  void test(List<String> list) {
+    long count = list.stream()
+      .ma<caret>p(e -> {
+        System.out.println(e);
+        // hello
+        return /* in return */ e;
+      })
+      .count();
+    System.out.println(count);
+  }
+}

java/java-tests/testData/inspection/streamApiCallChains/beforeMapToPeekPrimitive.java

@@ -0,0 +1,18 @@
+// "Replace with 'peek'" "true"
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
+
+public class Main {
+  void test() {
+    AtomicInteger counter = new AtomicInteger();
+    int[] ints = IntStream.range(0, 100)
+      .filter(x -> x % 3 == 0)
+      .m<caret>ap((x -> {
+        counter.incrementAndGet();
+        return x;
+      }))
+      .toArray();
+    System.out.println(counter.get());
+  }
+}

java/java-tests/testData/inspection/streamApiCallChains/beforeMapToPeekWritten.java

@@ -0,0 +1,16 @@
+// "Replace with 'peek'" "false"
+
+import java.util.stream.IntStream;
+
+public class Main {
+  void test() {
+    int[] ints = IntStream.range(0, 100)
+      .filter(x -> x % 3 == 0)
+      .m<caret>ap(x -> {
+        x++;
+        return x;
+      })
+      .toArray();
+    System.out.println(ints.length);
+  }
+}

resources-en/src/inspectionDescriptions/SimplifyStreamApiCallChains.html

@@ -15,12 +15,13 @@ It allows to avoid creating redundant temporary objects when traversing a collec
   <li><code>Collections.singleton().stream()</code> &rarr; <code>Stream.of()</code></li>
   <li><code>Collections.emptyList().stream()</code> &rarr; <code>Stream.empty()</code></li>
   <li><code>stream.filter().findFirst().isPresent()</code> &rarr; <code>stream.anyMatch()</code></li>
-  <li><code>stream.collect(Collectors.counting())</code> &rarr; <code>stream.count()</code></li>
-  <li><code>stream.collect(Collectors.maxBy())</code> &rarr; <code>stream.max()</code></li>
-  <li><code>stream.collect(Collectors.mapping())</code> &rarr; <code>stream.map().collect()</code></li>
-  <li><code>stream.collect(Collectors.reducing())</code> &rarr; <code>stream.reduce()</code></li>
-  <li><code>stream.collect(Collectors.summingInt())</code> &rarr; <code>stream.mapToInt().sum()</code></li>
+  <li><code>stream.collect(counting())</code> &rarr; <code>stream.count()</code></li>
+  <li><code>stream.collect(maxBy())</code> &rarr; <code>stream.max()</code></li>
+  <li><code>stream.collect(mapping())</code> &rarr; <code>stream.map().collect()</code></li>
+  <li><code>stream.collect(reducing())</code> &rarr; <code>stream.reduce()</code></li>
+  <li><code>stream.collect(summingInt())</code> &rarr; <code>stream.mapToInt().sum()</code></li>
   <li><code>stream.mapToObj(x -> x)</code> &rarr; <code>stream.boxed()</code></li>
+  <li><code>stream.map(x -> {...; return x;})</code> &rarr; <code>stream.peek(x -> ...)</code></li>
   <li><code>!stream.anyMatch()</code> &rarr; <code>stream.noneMatch()</code></li>
   <li><code>!stream.anyMatch(x -> !(...))</code> &rarr; <code>stream.allMatch()</code></li>
   <li><code>stream.map().anyMatch(Boolean::booleanValue)</code> -> <code>stream.anyMatch()</code></li>