StreamToLoop inspection: simplify toMap mergers (a,b)->a (use putIfAbsent) and (a,b)->b (use put)

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

@@ -29,6 +29,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -729,6 +730,27 @@ abstract class TerminalOperation extends Operation {
         return "if("+map+".put("+myKeyExtractor.getText()+","+myValueExtractor.getText()+")!=null) {\n"+
                "throw new java.lang.IllegalStateException(\"Duplicate key\");\n}\n";
       }
+      if(myMerger instanceof PsiLambdaExpression) {
+        PsiLambdaExpression lambda = (PsiLambdaExpression)myMerger;
+        PsiParameter[] parameters = lambda.getParameterList().getParameters();
+        if(parameters.length == 2) {
+          PsiExpression body = LambdaUtil.extractSingleExpressionFromBody(lambda.getBody());
+          if(body instanceof PsiReferenceExpression) {
+            PsiReferenceExpression ref = (PsiReferenceExpression)body;
+            if(ref.getQualifierExpression() == null) {
+              // cannot use isReferenceTo here as lambda could be detached from PsiFile
+              if (Objects.equals(parameters[0].getName(), ref.getReferenceName())) {
+                // like (a, b) -> a
+                return map + ".putIfAbsent(" + myKeyExtractor.getText() + "," + myValueExtractor.getText() + ");\n";
+              }
+              else if (Objects.equals(parameters[1].getName(), ref.getReferenceName())) {
+                // like (a, b) -> b
+                return map + ".put(" + myKeyExtractor.getText() + "," + myValueExtractor.getText() + ");\n";
+              }
+            }
+          }
+        }
+      }
       return map+".merge("+myKeyExtractor.getText()+","+myValueExtractor.getText()+","+myMerger.getText()+");\n";
     }
   }

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

@@ -0,0 +1,21 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  public static TreeMap<Integer, String> test(List<String> strings) {
+      TreeMap<Integer, String> map = new TreeMap<>();
+      for (String s1 : strings) {
+          if (!s1.isEmpty()) {
+              map.put(s1.length(), s1.trim());
+          }
+      }
+      return map;
+  }
+
+  public static void main(String[] args) {
+    System.out.println(test(Arrays.asList()));
+    System.out.println(test(Arrays.asList("a", "bbb", "cc", "d", "eee", "")));
+  }
+}

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

@@ -8,7 +8,7 @@ public class Main {
       TreeMap<Integer, String> map = new TreeMap<>();
       for (String s1 : strings) {
           if (!s1.isEmpty()) {
-              map.merge(s1.length(), s1, (s, string) -> s);
+              map.putIfAbsent(s1.length(), s1);
           }
       }
       return map;

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

@@ -0,0 +1,16 @@
+// "Replace Stream API chain with loop" "true"
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class Main {
+  public static TreeMap<Integer, String> test(List<String> strings) {
+    return strings.stream().filter(s -> !s.isEmpty())
+      .col<caret>lect(Collectors.toMap(String::length, String::trim, (s, string) -> string, TreeMap::new));
+  }
+
+  public static void main(String[] args) {
+    System.out.println(test(Arrays.asList()));
+    System.out.println(test(Arrays.asList("a", "bbb", "cc", "d", "eee", "")));
+  }
+}