From dced7d86d037a66016b07b1dd28c8be379e2a739 Mon Sep 17 00:00:00 2001
From: Mikhail Pyltsin <Mikhail.Pyltsin@jetbrains.com>
Date: Mon, 18 Sep 2023 10:43:45 +0200
Subject: [PATCH] [java-inspections] IDEA-332629 Support Optional.stream() in
 DFA

- added inlining for Optional.stream, based n SpecialField.OPTIONAL_VALUE

GitOrigin-RevId: 50287ae63b11083c133f9af54373310bf1b818f8
---
 .../java/inliner/StreamChainInliner.java      | 23 +++++++++++
 .../fixture/OptionalStreamInlining.java       | 41 +++++++++++++++++++
 .../DataFlowInspection9Test.java              |  2 +
 3 files changed, 66 insertions(+)
 create mode 100644 java/java-tests/testData/inspection/dataFlow/fixture/OptionalStreamInlining.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/inliner/StreamChainInliner.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/inliner/StreamChainInliner.java
index 031f59b33d56..1fd9628f5982 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/inliner/StreamChainInliner.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/inliner/StreamChainInliner.java
@@ -924,6 +924,29 @@ public class StreamChainInliner implements CallInliner {
         return;
       }
     }
+    if (OPTIONAL_STREAM.test(sourceCall)) {
+      PsiExpression qualifierExpression = sourceCall.getMethodExpression().getQualifierExpression();
+      if (qualifierExpression != null) {
+        PsiType optValueType = PsiUtil.substituteTypeParameter(qualifierExpression.getType(), JAVA_UTIL_OPTIONAL, 0, false);
+        if (optValueType != null) {
+          builder.pushExpression(qualifierExpression);
+          builder //optVar
+            .unwrap(SpecialField.OPTIONAL_VALUE) //optVar.value
+            .pushForWrite(builder.createTempVariable(qualifierExpression.getType()))//optVar.value optValueVar
+            .swap() //optValueVar optVar.value
+            .assign() //optValueVar
+            .dup() //optValueVar optValueVar
+            .chain(firstStep::before)
+            .ifNull() //optValueVar
+            // skip loop at all
+              .pop() //..
+            .elseBranch() //optValueVar
+              .chain(firstStep::iteration)
+            .end();
+          return;
+        }
+      }
+    }
     startStreamUnknown(builder, firstStep, originalQualifier, originalQualifierAlreadyChecked, inType);
   }
 
diff --git a/java/java-tests/testData/inspection/dataFlow/fixture/OptionalStreamInlining.java b/java/java-tests/testData/inspection/dataFlow/fixture/OptionalStreamInlining.java
new file mode 100644
index 000000000000..6a5853e94301
--- /dev/null
+++ b/java/java-tests/testData/inspection/dataFlow/fixture/OptionalStreamInlining.java
@@ -0,0 +1,41 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.*;
+import java.util.function.*;
+import java.util.stream.*;
+
+public class OptionalStreamInlining {
+  public static void testOptionalStream(String s) {
+    Optional<List<String>> strings = Optional.ofNullable(s)
+      .map(OptionalStreamInlining::getLetters);
+    Optional<String> s1 = strings
+      .stream()
+      .flatMap(Collection::stream)
+      .min(Comparator.naturalOrder());
+    if (s1.isPresent()) {
+      System.out.println(s.length());
+    } else {
+      System.out.println();
+      System.out.println(s.<warning descr="Method invocation 'length' may produce 'NullPointerException'">length</warning>());
+    }
+  }
+  public static void testOptionalStream2(String s) {
+    Optional<String> s1 = Optional.ofNullable(s)
+      .map(OptionalStreamInlining::getLetters)
+      .stream()
+      .flatMap(Collection::stream)
+      .min(Comparator.naturalOrder());
+    if (s1.isPresent()) {
+      System.out.println(s.length());
+    } else {
+      System.out.println();
+      System.out.println(s.<warning descr="Method invocation 'length' may produce 'NullPointerException'">length</warning>());
+    }
+  }
+  private static List<String> getLetters(String t){
+    ArrayList<String> strings = new ArrayList<>();
+    strings.add("1");
+    return strings;
+  }
+}
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java
index b2c79508b78a..9824178376aa 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java
@@ -22,4 +22,6 @@ public class DataFlowInspection9Test extends DataFlowInspectionTestCase {
   public void testMutabilityInferred() { doTest(); }
   public void testObjectsRequireNonNullElse() { doTest(); }
   public void testNewCollectionAliasing() { doTest(); }
+
+  public void testOptionalStreamInlining() { doTest(); }
 }
\ No newline at end of file