From 67ae5c45e2629eb74fd21c46882f44740a47c4b7 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Fri, 25 Feb 2022 12:38:10 +0700
Subject: [PATCH] [java-dfa] Report when nullable value is passed to
 findFirst()

Fixes IDEA-289362 NPE could be detected on findFirst on a stream with nullable elements

GitOrigin-RevId: 2d27d29eacdb2612608060d423646184a9093f51
---
 .../java/inliner/StreamChainInliner.java       | 11 +++++++++++
 .../fixture/StreamFindFirstExpectNotNull.java  | 18 ++++++++++++++++++
 .../dataFlow/fixture/StreamInlining.java       |  2 +-
 .../DataFlowInspection8Test.java               |  1 +
 4 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 java/java-tests/testData/inspection/dataFlow/fixture/StreamFindFirstExpectNotNull.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 9cc4d6e09e61..d9cc535b94b4 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
@@ -382,8 +382,14 @@ public class StreamChainInliner implements CallInliner {
   }
 
   static class OptionalTerminalStep extends TerminalStep {
+    private final boolean myExpectNotNull;
+
     OptionalTerminalStep(@NotNull PsiMethodCallExpression call) {
       super(call, ArrayUtil.getFirstElement(call.getArgumentList().getExpressions()));
+      String methodName = call.getMethodExpression().getReferenceName();
+      // While findFirst/findAny will work if the found element is definitely not-null,
+      // it's highly suspicious to supply a nullable element there
+      myExpectNotNull = "findFirst".equals(methodName) || "findAny".equals(methodName);
     }
 
     @Override
@@ -405,6 +411,11 @@ public class StreamChainInliner implements CallInliner {
       DfType source = DfaOptionalSupport.getOptionalValue(true);
       builder.pushForWrite(myResult).push(source).assign().splice(2);
     }
+
+    @Override
+    boolean expectNotNull() {
+      return myExpectNotNull;
+    }
   }
 
   static class MinMaxTerminalStep extends TerminalStep {
diff --git a/java/java-tests/testData/inspection/dataFlow/fixture/StreamFindFirstExpectNotNull.java b/java/java-tests/testData/inspection/dataFlow/fixture/StreamFindFirstExpectNotNull.java
new file mode 100644
index 000000000000..90cb81434809
--- /dev/null
+++ b/java/java-tests/testData/inspection/dataFlow/fixture/StreamFindFirstExpectNotNull.java
@@ -0,0 +1,18 @@
+import org.jetbrains.annotations.*;
+
+import java.util.*;
+import java.util.function.*;
+import java.util.stream.*;
+
+class StreamTypeAnnoInlining {
+  @Nullable
+  private static Object fun(String s) {
+    return null;
+  }
+
+  void test() {
+    Stream.of("a", "b").map(s -> <warning descr="Function may return null, but it's not allowed here">fun(s)</warning>).findFirst();
+    Stream.of("a", "b").map(s -> <warning descr="Function may return null, but it's not allowed here">fun(s)</warning>).findAny();
+  }
+
+}
diff --git a/java/java-tests/testData/inspection/dataFlow/fixture/StreamInlining.java b/java/java-tests/testData/inspection/dataFlow/fixture/StreamInlining.java
index 8b5bf7e6ef0d..69fabea4e784 100644
--- a/java/java-tests/testData/inspection/dataFlow/fixture/StreamInlining.java
+++ b/java/java-tests/testData/inspection/dataFlow/fixture/StreamInlining.java
@@ -133,7 +133,7 @@ public class StreamInlining {
 
   String blockLambda(List<String> list) {
     return list == null ? "" : list.stream().map(s -> {
-      return s.equals("abc") ? null : s;
+      return s.equals("abc") ? <warning descr="Function may return null, but it's not allowed here">null</warning> : s;
     }).findFirst().orElse("");
   }
 
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
index 2efa0ffaabde..14258f20134c 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection8Test.java
@@ -206,6 +206,7 @@ public class DataFlowInspection8Test extends DataFlowInspectionTestCase {
     setupTypeUseAnnotations("foo", myFixture);
     doTest();
   }
+  public void testStreamFindFirstExpectNotNull() { doTest(); }
   public void testStreamAnyMatchIsNull() { doTest(); }
   public void testStreamCustomSumMethod() { doTest(); }
   public void testStreamReduceLogicalAnd() { doTest(); }