StreamApiMigrationInspection: enable flatMap with primitive type change

java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java

@@ -837,17 +837,22 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
     @Override
     public String createReplacement() {
       String operation = "flatMap";
-      PsiType type = mySource.getVariable().getType();
-      if(type instanceof PsiPrimitiveType && !type.equals(myVariable.getType())) {
-        if(type.equals(PsiType.INT)) {
+      PsiType inType = myVariable.getType();
+      PsiType outType = mySource.getVariable().getType();
+      String lambda = myVariable.getName() + " -> " + getStreamExpression();
+      if(outType instanceof PsiPrimitiveType && !outType.equals(inType)) {
+        if(outType.equals(PsiType.INT)) {
           operation = "flatMapToInt";
-        } else if(type.equals(PsiType.LONG)) {
+        } else if(outType.equals(PsiType.LONG)) {
           operation = "flatMapToLong";
-        } else if(type.equals(PsiType.DOUBLE)) {
+        } else if(outType.equals(PsiType.DOUBLE)) {
           operation = "flatMapToDouble";
         }
       }
-      return "." + operation + "(" + myVariable.getName() + " -> " + getStreamExpression() + ")";
+      if(inType instanceof PsiPrimitiveType && !outType.equals(inType)) {
+        return ".mapToObj(" + lambda + ")." + operation + "(" + CommonClassNames.JAVA_UTIL_FUNCTION_FUNCTION + ".identity())";
+      }
+      return "." + operation + "(" + lambda + ")";
     }
 
     @NotNull

java/java-impl/src/com/intellij/codeInspection/streamMigration/TerminalBlock.java

@@ -163,10 +163,6 @@ class TerminalBlock {
       StreamSource source = StreamSource.tryCreate(loopStatement);
       final PsiStatement body = loopStatement.getBody();
       if(source == null || body == null) return null;
-      // flatMap from primitive to primitive is supported only if primitive types match
-      // otherwise it would be necessary to create bogus step like
-      // .mapToObj(var -> collection.stream()).flatMap(Function.identity())
-      if(myVariable.getType() instanceof PsiPrimitiveType && !myVariable.getType().equals(source.getVariable().getType())) return null;
       FlatMapOp op = new FlatMapOp(source, myVariable);
       TerminalBlock withFlatMap = new TerminalBlock(this, op, source.getVariable(), body);
       if(!VariableAccessUtils.variableIsUsed(myVariable, body)) {

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFlatMapChangePrimitiveType.java

@@ -0,0 +1,13 @@
+// "Replace with count()" "true"
+import java.util.Arrays;
+import java.util.function.Function;
+import java.util.stream.IntStream;
+import java.util.stream.LongStream;
+import java.util.stream.Stream;
+
+public class Test {
+  public void test() {
+      long count = IntStream.range(0, 10).mapToObj(i -> LongStream.range(0, i)).flatMapToLong(Function.identity()).mapToObj(l -> Stream.of("x", "y", "z")).flatMap(Function.identity()).flatMapToInt(s -> IntStream.range(0, s.length())).count();
+      System.out.println(count);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFlatMapChangePrimitiveType.java

@@ -0,0 +1,18 @@
+// "Replace with count()" "true"
+import java.util.Arrays;
+
+public class Test {
+  public void test() {
+    long count = 0;
+    for(int <caret>i=0; i<10; i++) {
+      for(long l = 0; l < i; l++) {
+        for(String s : Arrays.asList("x", "y", "z")) {
+          for(int k = 0; k < s.length(); k++) {
+            count++;
+          }
+        }
+      }
+    }
+    System.out.println(count);
+  }
+}