[java-inspections] Hard-code Map.compute second lambda parameter as nullable

Fixes IDEA-290431 False positive for “@NotNull/@Nullable problems” inspection on Map.compute

GitOrigin-RevId: fdac0a5eaa0aca3bdd1bef3d9c2eb4c445f7e66e

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DfaPsiUtil.java

@@ -279,16 +279,21 @@ public final class DfaPsiUtil {
     return Nullability.UNKNOWN;
   }
 
+  private static final CallMatcher OPTIONAL_FUNCTIONS =
+    CallMatcher.instanceCall(JAVA_UTIL_OPTIONAL, "map", "filter", "ifPresent", "flatMap", "ifPresentOrElse");
+  private static final CallMatcher MAP_COMPUTE =
+    CallMatcher.instanceCall(JAVA_UTIL_MAP, "compute").parameterTypes("K", JAVA_UTIL_FUNCTION_BI_FUNCTION);
+
   @NotNull
   private static Nullability getLambdaParameterNullability(@NotNull PsiMethod method, int parameterIndex, int lambdaParameterIndex) {
-    PsiClass type = method.getContainingClass();
-    if(type != null) {
-      if(JAVA_UTIL_OPTIONAL.equals(type.getQualifiedName())) {
-        String methodName = method.getName();
-        if((methodName.equals("map") || methodName.equals("filter") || methodName.equals("ifPresent") || methodName.equals("flatMap"))
-          && parameterIndex == 0 && lambdaParameterIndex == 0) {
-          return Nullability.NOT_NULL;
-        }
+    if (OPTIONAL_FUNCTIONS.methodMatches(method)) {
+      if (parameterIndex == 0 && lambdaParameterIndex == 0) {
+        return Nullability.NOT_NULL;
+      }
+    }
+    else if (MAP_COMPUTE.methodMatches(method)) {
+      if (parameterIndex == 1 && lambdaParameterIndex == 1) {
+        return Nullability.NULLABLE;
       }
     }
     return Nullability.UNKNOWN;

java/java-tests/testData/inspection/nullableProblems/MapComputeLambdaAnnotation.java

@@ -0,0 +1,17 @@
+import typeUse.NotNull;
+import typeUse.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class MapComputeLambdaAnnotation {
+  public static void main(final String[] args) {
+    final Map<String, @NotNull String> test = new HashMap<>();
+
+    test.compute("first", (String a, @Nullable String b) -> {
+      assert b == null;
+
+      return null;
+    });
+  }
+}

java/java-tests/testSrc/com/intellij/java/codeInspection/NullableStuffInspectionTest.java

@@ -405,4 +405,9 @@ public class NullableStuffInspectionTest extends LightJavaCodeInsightFixtureTest
                        "}");
     doTest();
   }
+
+  public void testMapComputeLambdaAnnotation() {
+    DataFlowInspection8Test.setupTypeUseAnnotations("typeUse", myFixture);
+    doTest();
+  }
 }