IDEA-CR-14049 missing getType() added; support array initializers in Optional map

java/java-analysis-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithFindFirstFix.java

@@ -55,7 +55,7 @@ class ReplaceWithFindFirstFix extends MigrateToStreamFix {
       if (nextReturnStatement == null) return;
       PsiExpression orElseExpression = nextReturnStatement.getReturnValue();
       if (!ExpressionUtils.isSimpleExpression(orElseExpression)) return;
-      stream = generateOptionalUnwrap(stream, tb, value, orElseExpression);
+      stream = generateOptionalUnwrap(stream, tb, value, orElseExpression, null);
       restoreComments(foreachStatement, body);
       boolean siblings = nextReturnStatement.getParent() == foreachStatement.getParent();
       PsiElement result = foreachStatement.replace(elementFactory.createStatementFromText("return " + stream + ";", foreachStatement));
@@ -81,19 +81,19 @@ class ReplaceWithFindFirstFix extends MigrateToStreamFix {
       if (status != InitializerUsageStatus.UNKNOWN) {
         PsiExpression initializer = var.getInitializer();
         if (initializer != null) {
-          String replacementText = generateOptionalUnwrap(stream, tb, value, initializer);
+          String replacementText = generateOptionalUnwrap(stream, tb, value, initializer, var.getType());
           replaceInitializer(foreachStatement, var, initializer, replacementText, status);
           return;
         }
       }
       PsiElement result = foreachStatement.replace(elementFactory.createStatementFromText(
-        var.getName() + " = " + generateOptionalUnwrap(stream, tb, value, lValue) + ";", foreachStatement));
+        var.getName() + " = " + generateOptionalUnwrap(stream, tb, value, lValue, var.getType()) + ";", foreachStatement));
       simplifyAndFormat(project, result);
     }
   }
 
   private static String generateOptionalUnwrap(String stream, @NotNull StreamApiMigrationInspection.TerminalBlock tb,
-                                               PsiExpression trueExpression, PsiExpression falseExpression) {
+                                               PsiExpression trueExpression, PsiExpression falseExpression, PsiType targetType) {
     PsiVariable var = tb.getVariable();
     if (!StreamApiMigrationInspection.isIdentityMapping(var, trueExpression)) {
       if(trueExpression instanceof PsiTypeCastExpression && ExpressionUtils.isNullLiteral(falseExpression)) {
@@ -112,9 +112,11 @@ class ReplaceWithFindFirstFix extends MigrateToStreamFix {
         if(EquivalenceChecker.getCanonicalPsiEquivalence().expressionsAreEquivalent(falseExpression, condition.getElseExpression())) {
           return generateOptionalUnwrap(
             stream + ".filter(" + LambdaUtil.createLambda(var, condition.getCondition()) + ")", tb,
-            condition.getThenExpression(), falseExpression);
+            condition.getThenExpression(), falseExpression, var.getType());
         }
       }
+      trueExpression =
+        targetType == null ? trueExpression : ExpressionUtils.convertInitializerToNormalExpression(trueExpression, targetType);
       stream += ".map(" + LambdaUtil.createLambda(var, trueExpression) + ")";
     }
     stream += ".orElse(" + falseExpression.getText() + ")";

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

@@ -586,7 +586,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
             if(!(var instanceof PsiVariable) || !nonFinalVariables.contains(var)) return;
             PsiExpression rValue = assignment.getRExpression();
             if(rValue == null || isVariableReferenced((PsiVariable)var, rValue)) return;
-            if(tb.getVariable() instanceof PsiPrimitiveType && !isIdentityMapping(tb.getVariable(), rValue)) return;
+            if(tb.getVariable().getType() instanceof PsiPrimitiveType && !isIdentityMapping(tb.getVariable(), rValue)) return;
             registerProblem(statement, "findFirst", new ReplaceWithFindFirstFix());
           }
         }

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

@@ -0,0 +1,12 @@
+// "Replace with findFirst()" "true"
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public class Main {
+  public int[] testMap(Map<String, List<String>> map) throws Exception {
+      int[] arr = map.values().stream().filter(Objects::nonNull).findFirst().map(list -> new int[]{list.size()}).orElse(null);
+      return arr;
+  }
+}

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

@@ -0,0 +1,11 @@
+// "Replace with findFirst()" "true"
+
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  public void testMap(Map<String, List<String>> map) throws Exception {
+      int bigSize = map.values().stream().mapToInt(List::size).filter(size -> size > 10).findFirst().orElse(0);
+      System.out.println(bigSize);
+  }
+}

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

@@ -0,0 +1,17 @@
+// "Replace with findFirst()" "true"
+
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  public int[] testMap(Map<String, List<String>> map) throws Exception {
+    int[] arr = null;
+    for(List<String> list : map.valu<caret>es()) {
+      if(list != null) {
+        arr = {list.size()};
+        break;
+      }
+    }
+    return arr;
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with findFirst()" "true"
+
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  public void testMap(Map<String, List<String>> map) throws Exception {
+    int bigSize = 0;
+    for(List<String> list : map.valu<caret>es()) {
+      int size = list.size();
+      if(size > 10) {
+        bigSize = size;
+        break;
+      }
+    }
+    System.out.println(bigSize);
+  }
+}

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

@@ -0,0 +1,18 @@
+// "Replace with findFirst()" "false"
+
+import java.util.List;
+import java.util.Map;
+
+public class Main {
+  public void testMap(Map<String, List<String>> map) throws Exception {
+    int bigSize = 0;
+    for(List<String> list : map.valu<caret>es()) {
+      int size = list.size();
+      if(size > 10) {
+        bigSize = size*2;
+        break;
+      }
+    }
+    System.out.println(bigSize);
+  }
+}