From ce18179a789da6bc970e2e25bb235ca4279f0c5d Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Fri, 6 Dec 2024 09:55:28 +0100
Subject: [PATCH] [java-analysis] Parameter nullability: prefer nullability
 known from type over nullability known from parameter declaration

Type nullability could be more precise if parameter is generic
Fixes IDEA-364343 False-positive NPE at unboxing inside lambda with JSpecify annotations

GitOrigin-RevId: 9a49f5687eccaa013e639cdf15950be911e100bc
---
 .../codeInspection/dataFlow/DfaPsiUtil.java   | 10 ++++-----
 .../fixture/JSpecifyUnboxingInLambda.java     | 21 +++++++++++++++++++
 .../DataFlowInspection9Test.java              |  6 ++++++
 3 files changed, 32 insertions(+), 5 deletions(-)
 create mode 100644 java/java-tests/testData/inspection/dataFlow/fixture/JSpecifyUnboxingInLambda.java

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaPsiUtil.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaPsiUtil.java
index d7de18fb510d2..e9eda28a53bca 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaPsiUtil.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaPsiUtil.java
@@ -307,12 +307,12 @@ public static Nullability getFunctionalParameterNullability(PsiFunctionalExpress
     if (sam != null) {
       PsiParameter parameter = sam.getParameterList().getParameter(index);
       if (parameter != null) {
-        Nullability nullability = getElementNullability(null, parameter);
-        if (nullability != Nullability.UNKNOWN) {
-          return nullability;
-        }
         PsiType parameterType = type.resolveGenerics().getSubstitutor().substitute(parameter.getType());
-        return getTypeNullability(GenericsUtil.eliminateWildcards(parameterType, false, true));
+        NullabilityAnnotationInfo info = getTypeNullabilityInfo(GenericsUtil.eliminateWildcards(parameterType, false, true));
+        if (info != null) {
+          return info.getNullability();
+        }
+        return getElementNullability(null, parameter);
       }
     }
     return Nullability.UNKNOWN;
diff --git a/java/java-tests/testData/inspection/dataFlow/fixture/JSpecifyUnboxingInLambda.java b/java/java-tests/testData/inspection/dataFlow/fixture/JSpecifyUnboxingInLambda.java
new file mode 100644
index 0000000000000..3dfb1a28d52d5
--- /dev/null
+++ b/java/java-tests/testData/inspection/dataFlow/fixture/JSpecifyUnboxingInLambda.java
@@ -0,0 +1,21 @@
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+// IDEA-364343
+@NullMarked
+class AnotherActivity {
+  public interface ThrowingFunction<T1 extends @Nullable Object, T2 extends @Nullable Object> {
+    T2 apply(T1 input) throws Throwable;
+  }
+
+  abstract static class Decoder<T extends @Nullable Object> {
+    abstract <T2 extends @Nullable Object> Decoder<T2> then(
+      ThrowingFunction<? super T, ? extends T2> dataTransform);
+  }
+
+  native Decoder<Boolean> foo();
+
+  Decoder<Boolean> doWork() {
+    return foo().then(f -> !f);
+  }
+}
\ No newline at end of file
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 c9f1e21eca750..260710253b8df 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/DataFlowInspection9Test.java
@@ -46,4 +46,10 @@ public void testJSpecifyUpperBound() {
     setupTypeUseAnnotations("org.jspecify.annotations", myFixture);
     doTest();
   }
+
+  public void testJSpecifyUnboxingInLambda() {
+    addJSpecifyNullMarked(myFixture);
+    setupTypeUseAnnotations("org.jspecify.annotations", myFixture);
+    doTest();
+  }
 }
\ No newline at end of file