IDEA-202753 Map.entrySet can be simplified analysis could identify more cases.

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

@@ -1917,9 +1917,13 @@ public class SimplifyStreamApiCallChainsInspection extends AbstractBaseJavaLocal
 
   static class EntrySetMapFix implements CallChainSimplification {
     private final String myMapMethod;
+    private final boolean myDeleteMap;
+    private final String[] myNames;
 
-    EntrySetMapFix(String mapMethod) {
-      myMapMethod = mapMethod;
+    EntrySetMapFix(String entryMethod, boolean deleteMap) {
+      myMapMethod = entryMethod.equals("getKey") ? "keySet" : "values";
+      myDeleteMap = deleteMap;
+      myNames = myMapMethod.equals("keySet") ? new String[]{"k", "key"} : new String[]{"v", "value"};
     }
 
     @Override
@@ -1936,30 +1940,116 @@ public class SimplifyStreamApiCallChainsInspection extends AbstractBaseJavaLocal
     public PsiElement simplify(PsiMethodCallExpression call) {
       PsiMethodCallExpression qualifierCall = getQualifierMethodCall(call);
       if (qualifierCall == null) return null;
-      PsiMethodCallExpression qualifierQualifierCall = getQualifierMethodCall(qualifierCall);
-      if (qualifierQualifierCall == null) return null;
-      CommentTracker ct = new CommentTracker();
-      PsiMethodCallExpression result = (PsiMethodCallExpression)ct.replaceAndRestoreComments(call, qualifierCall);
-      PsiMethodCallExpression newQualifier = Objects.requireNonNull(getQualifierMethodCall(result));
-      ExpressionUtils.bindCallTo(newQualifier, myMapMethod);
+      if (myDeleteMap) {
+        CommentTracker ct = new CommentTracker();
+        call = (PsiMethodCallExpression)ct.replaceAndRestoreComments(call, qualifierCall);
+      }
+      PsiMethodCallExpression result = call;
+      while (call != null) {
+        if (STREAM_MAP_TO_ALL.test(call) || STREAM_FILTER.test(call)) {
+          PsiExpression arg = PsiUtil.skipParenthesizedExprDown(call.getArgumentList().getExpressions()[0]);
+          if (arg instanceof PsiLambdaExpression) {
+            updateLambda((PsiLambdaExpression)arg);
+          }
+          else if (arg instanceof PsiMethodReferenceExpression) {
+            PsiType type = LambdaUtil.getFunctionalInterfaceReturnType((PsiFunctionalExpression)arg);
+            String name = new VariableNameGenerator(arg, VariableKind.PARAMETER).byType(type).byName(myNames).generate(false);
+            new CommentTracker().replaceAndRestoreComments(arg, name + "->" + name);
+          }
+        }
+        call = getQualifierMethodCall(call);
+        if (MAP_ENTRY_SET.test(call)) {
+          ExpressionUtils.bindCallTo(call, myMapMethod);
+          break;
+        }
+      }
+      LambdaCanBeMethodReferenceInspection.replaceAllLambdasWithMethodReferences(result);
       return result;
     }
 
+    private void updateLambda(PsiLambdaExpression lambda) {
+      PsiParameter[] parameters = lambda.getParameterList().getParameters();
+      if (parameters.length != 1) return;
+      PsiParameter parameter = parameters[0];
+      PsiType type = PsiUtil.substituteTypeParameter(parameter.getType(), JAVA_UTIL_MAP_ENTRY, 1, true);
+      PsiElement body = lambda.getBody();
+      if (body == null) return;
+      List<PsiMethodCallExpression> calls = new ArrayList<>();
+      PsiLocalVariable declaration = null;
+      for (PsiReferenceExpression ref : VariableAccessUtils.getVariableReferences(parameter, body)) {
+        PsiMethodCallExpression call = ExpressionUtils.getCallForQualifier(ref);
+        if (call != null) {
+          calls.add(call);
+          if (declaration == null) {
+            PsiLocalVariable var = tryCast(PsiUtil.skipParenthesizedExprUp(call.getParent()), PsiLocalVariable.class);
+            if (var != null && var.getParent() instanceof PsiDeclarationStatement && var.getParent().getParent() == body) {
+              declaration = var;
+            }
+          }
+        }
+      }
+      String name = declaration == null ? null : declaration.getName();
+      if (name == null) {
+        name = new VariableNameGenerator(lambda, VariableKind.PARAMETER).byType(type).byName(myNames).generate(false);
+      }
+      for (PsiMethodCallExpression call : calls) {
+        PsiElement result = new CommentTracker().replaceAndRestoreComments(call, name);
+        if (call == body) {
+          body = result;
+        }
+      }
+      if (declaration != null) {
+        new CommentTracker().deleteAndRestoreComments(declaration);
+      }
+      CommentTracker ct = new CommentTracker();
+      ct.replaceAndRestoreComments(lambda, name + "->" + ct.text(body));
+    }
+
     public static CallHandler<CallChainSimplification> handler() {
-      return CallHandler.of(STREAM_MAP, call -> {
+      return CallHandler.of(STREAM_MAP_TO_ALL, call -> {
         PsiMethodCallExpression qualifierCall = getQualifierMethodCall(call);
+        List<String> methods = new ArrayList<>(Arrays.asList("getKey", "getValue"));
+        while (STREAM_FILTER.test(qualifierCall)) {
+          String methodName = getSingleCalledMethodName(qualifierCall.getArgumentList().getExpressions()[0]);
+          if (methodName == null || !methods.contains(methodName)) return null;
+          methods = Collections.singletonList(methodName);
+          qualifierCall = getQualifierMethodCall(qualifierCall);
+        }
         if (!COLLECTION_STREAM.test(qualifierCall)) return null;
         PsiMethodCallExpression qualifierQualifierCall = getQualifierMethodCall(qualifierCall);
         if (!MAP_ENTRY_SET.test(qualifierQualifierCall)) return null;
         PsiExpression arg = call.getArgumentList().getExpressions()[0];
-        if (FunctionalExpressionUtils.isFunctionalReferenceTo(arg, JAVA_UTIL_MAP_ENTRY, null, "getKey")) {
-          return new EntrySetMapFix("keySet");
+        for (String method : methods) {
+          if (FunctionalExpressionUtils.isFunctionalReferenceTo(arg, JAVA_UTIL_MAP_ENTRY, null, method)) {
+            return new EntrySetMapFix(method, "map".equals(call.getMethodExpression().getReferenceName()));
+          }
         }
-        if (FunctionalExpressionUtils.isFunctionalReferenceTo(arg, JAVA_UTIL_MAP_ENTRY, null, "getValue")) {
-          return new EntrySetMapFix("values");
+        String methodName = getSingleCalledMethodName(arg);
+        if (methodName != null && methods.contains(methodName)) {
+          return new EntrySetMapFix(methodName, false);
         }
         return null;
       });
     }
+
+    @Nullable
+    private static String getSingleCalledMethodName(PsiExpression arg) {
+      PsiLambdaExpression lambda = tryCast(PsiUtil.skipParenthesizedExprDown(arg), PsiLambdaExpression.class);
+      if (lambda == null) return null;
+      PsiParameter[] parameters = lambda.getParameterList().getParameters();
+      if (parameters.length != 1) return null;
+      PsiParameter parameter = parameters[0];
+      PsiElement body = lambda.getBody();
+      if (body == null) return null;
+      String methodName = null;
+      for (PsiReferenceExpression ref : VariableAccessUtils.getVariableReferences(parameter, body)) {
+        PsiMethodCallExpression call = ExpressionUtils.getCallForQualifier(ref);
+        if (call == null || !call.getArgumentList().isEmpty()) return null;
+        String name = call.getMethodExpression().getReferenceName();
+        if (name == null || (methodName != null && !methodName.equals(name))) return null;
+        methodName = name;
+      }
+      return methodName;
+    }
   }
 }

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

+// "Fix all 'Stream API call chain can be simplified' problems in file" "true"
+import java.util.Map;
+import java.util.stream.Collectors;
+
+class Scratch {
+  void testSimple(Map<String, String> map) {
+    String res = map.values().stream().map(String::trim).collect(Collectors.joining(";"))
+    String res2 = map.keySet().stream().map(String::trim).collect(Collectors.joining(";"))
+    String res3 = map.entrySet().stream().map(e -> e.toString().trim()).collect(Collectors.joining(";"))
+  }
+
+  void testFilter(Map<String, String> map) {
+    String res = map.values().stream().filter(s -> !s.isEmpty()).map(String::trim)
+      .collect(Collectors.joining(";"))
+    String res2 = map.entrySet().stream().filter(e -> !e.getValue().isEmpty()).map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+    String res3 = map.keySet().stream().filter(s -> !s.isEmpty()).map(String::trim)
+      .collect(Collectors.joining(";"))
+  }
+
+  void testFilter2(Map<String, String> map) {
+    String res = map.values().stream().filter(s -> !s.isEmpty())
+      .filter(s -> s.startsWith("foo"))
+      .map(String::trim)
+      .collect(Collectors.joining(";"))
+    String res2 = map.entrySet().stream().filter(e -> !e.getValue().isEmpty())
+      .filter(e -> e.getKey().startsWith("foo"))
+      .map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+    String res3 = map.keySet().stream().filter(s -> !s.isEmpty())
+      .filter(s -> s.startsWith("foo"))
+      .map(String::trim)
+      .collect(Collectors.joining(";"))
+  }
+
+  void testDecl(Map<String, String> map) {
+      /*5*//*6*/
+      String res = map.values(/*0*/).stream().filter(val1 -> {
+          /*1*/
+          return !val1./*2*/isEmpty();
+      }).filter(val2 -> {
+          return val2.startsWith("foo");
+      }).map(/*3*//*4*/String::trim).collect(Collectors.joining(";"))
+
+  }
+
+  void testUnboxing(Map<String, Integer> map) {
+    int sum = map.values().stream().filter(integer -> integer > 0)
+      .mapToInt(integer -> integer + 1).sum();
+    int sum2 = map.values().stream().filter(integer -> integer > 0)
+      .mapToInt(integer -> integer).sum();
+    int sum3 = map.values().stream().filter(integer -> integer > 0)
+      .mapToInt(i -> i).sum();
+  }
+}
\ No newline at end of file

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

+// "Fix all 'Stream API call chain can be simplified' problems in file" "true"
+import java.util.Map;
+import java.util.stream.Collectors;
+
+class Scratch {
+  void testSimple(Map<String, String> map) {
+    String res = map.entrySet().stream().m<caret>ap(e -> e.getValue().trim()).collect(Collectors.joining(";"))
+    String res2 = map.entrySet().stream().map(e -> e.getKey().trim()).collect(Collectors.joining(";"))
+    String res3 = map.entrySet().stream().map(e -> e.toString().trim()).collect(Collectors.joining(";"))
+  }
+
+  void testFilter(Map<String, String> map) {
+    String res = map.entrySet().stream().filter(e -> !e.getValue().isEmpty()).map(e -> e.getValue().trim())
+      .collect(Collectors.joining(";"))
+    String res2 = map.entrySet().stream().filter(e -> !e.getValue().isEmpty()).map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+    String res3 = map.entrySet().stream().filter(e -> !e.getKey().isEmpty()).map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+  }
+
+  void testFilter2(Map<String, String> map) {
+    String res = map.entrySet().stream().filter(e -> !e.getValue().isEmpty())
+      .filter(e -> e.getValue().startsWith("foo"))
+      .map(e -> e.getValue().trim())
+      .collect(Collectors.joining(";"))
+    String res2 = map.entrySet().stream().filter(e -> !e.getValue().isEmpty())
+      .filter(e -> e.getKey().startsWith("foo"))
+      .map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+    String res3 = map.entrySet().stream().filter(e -> !e.getKey().isEmpty())
+      .filter(e -> e.getKey().startsWith("foo"))
+      .map(e -> e.getKey().trim())
+      .collect(Collectors.joining(";"))
+  }
+
+  void testDecl(Map<String, String> map) {
+    String res = map.entrySet(/*0*/).stream().filter(e -> {
+      String val1 = e./*1*/getValue();
+      return !val1./*2*/isEmpty();
+    }).filter(e -> {
+      String val2 = e.getValue();
+      return val2.startsWith("foo");
+    }).map((Entry<String, String> /*3*/e/*4*/) -> {
+      String val3 = e.getValue()/*5*/;
+      return val3.trim(/*6*/);
+    }).collect(Collectors.joining(";"))
+
+  }
+
+  void testUnboxing(Map<String, Integer> map) {
+    int sum = map.entrySet().stream().filter(x -> x.getValue() > 0)
+      .mapToInt(x -> x.getValue() + 1).sum();
+    int sum2 = map.entrySet().stream().filter(x -> x.getValue() > 0)
+      .mapToInt(x -> x.getValue()).sum();
+    int sum3 = map.entrySet().stream().filter(x -> x.getValue() > 0)
+      .mapToInt(Map.Entry::getValue).sum();
+  }
+}
\ No newline at end of file