[java-dfa] Support groupingBy/partitioningBy/groupingByConcurrent collectors, to some extent (w/o downstream)

Fixes IDEA-302543 'Optional.get()' without 'isPresent()' in groupingBy Collector despite checking it earlier with filter in the stream

GitOrigin-RevId: c562ca1479e32c16517474db937528808c1f0fb8

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/inliner/StreamChainInliner.java

@@ -84,6 +84,8 @@ public class StreamChainInliner implements CallInliner {
           staticCall(JAVA_UTIL_STREAM_COLLECTORS, "toCollection").parameterCount(1));
   private static final CallMatcher MAP_COLLECTOR =
     staticCall(JAVA_UTIL_STREAM_COLLECTORS, "toMap", "toConcurrentMap", "toUnmodifiableMap");
+  private static final CallMatcher GROUPING_COLLECTOR =
+    staticCall(JAVA_UTIL_STREAM_COLLECTORS, "groupingBy", "partitioningBy", "groupingByConcurrent");
   private static final CallMatcher NOT_NULL_COLLECTORS =
     staticCall(JAVA_UTIL_STREAM_COLLECTORS, "joining", "maxBy", "minBy", "averagingInt", "averagingLong", "averagingDouble",
                "summingInt", "summingLong", "summingDouble", "summarizingInt", "summarizingLong", "summarizingDouble");
@@ -689,6 +691,9 @@ public class StreamChainInliner implements CallInliner {
         } else {
           dfType = dfType.meet(DfTypes.LOCAL_OBJECT);
         }
+        if (SpecialField.fromQualifierType(dfType) == SpecialField.COLLECTION_SIZE) {
+          dfType = dfType.meet(SpecialField.COLLECTION_SIZE.asDfType(DfTypes.intValue(0)));
+        }
         builder.push(dfType);
       }
     }
@@ -779,6 +784,35 @@ public class StreamChainInliner implements CallInliner {
     }
   }
 
+  static class GroupingStep extends AbstractCollectionStep {
+    private final @NotNull PsiExpression myKeyExtractor;
+    private final @Nullable PsiExpression myDownstream;
+
+    GroupingStep(@NotNull PsiMethodCallExpression call,
+                 @NotNull PsiExpression keyExtractor,
+                 @Nullable PsiExpression downstream,
+                 @Nullable PsiExpression supplier) {
+      super(call, supplier, false);
+      myKeyExtractor = keyExtractor;
+      myDownstream = downstream;
+    }
+
+    @Override
+    void before(CFGBuilder builder) {
+      builder.evaluateFunction(myKeyExtractor)
+             .evaluateFunction(myDownstream);
+      super.before(builder);
+    }
+
+    @Override
+    void iteration(CFGBuilder builder) {
+      // Null keys are not tolerated
+      builder.invokeFunction(1, myKeyExtractor, Nullability.NOT_NULL);
+      // Actual addition of Map element is unnecessary for current analysis
+      builder.flush(SpecialField.COLLECTION_SIZE.createValue(builder.getFactory(), myResult)).pop();
+    }
+  }
+
   @Override
   public boolean tryInlineCall(@NotNull CFGBuilder builder, @NotNull PsiMethodCallExpression call) {
     if (TERMINAL_CALL.test(call)) {
@@ -948,6 +982,15 @@ public class StreamChainInliner implements CallInliner {
                              "toUnmodifiableMap".equals(collectorCall.getMethodExpression().getReferenceName()));
       }
     }
+    if (GROUPING_COLLECTOR.matches(collectorCall)) {
+      PsiExpression[] args = collectorCall.getArgumentList().getExpressions();
+      if (args.length >= 1 && args.length <= 3) {
+        PsiExpression keyExtractor = args[0];
+        PsiExpression downstream = args.length > 1 ? args[args.length - 1] : null;
+        PsiExpression supplier = args.length == 3 ? args[1] : null;
+        return new GroupingStep(call, keyExtractor, downstream, supplier);
+      }
+    }
 
     return new UnknownTerminalStep(call, NOT_NULL_COLLECTORS.test(collectorCall));
   }

java/java-tests/testData/inspection/dataFlow/fixture/StreamGroupingBy.java

@@ -0,0 +1,31 @@
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+class Demo {
+  void test(List<String> list) {
+    Map<String, List<String>> map1 = list.stream()
+      .filter(Objects::isNull)
+      .collect(Collectors.groupingBy(x -> x.<warning descr="Method invocation 'trim' will produce 'NullPointerException'">trim</warning>()));
+    if (list.isEmpty() && <warning descr="Condition 'map1.isEmpty()' is always 'true' when reached">map1.isEmpty()</warning>) {}
+    Map<String, List<String>> map2 = list.stream()
+      .filter(Objects::isNull)
+      .collect(Collectors.groupingByConcurrent(x -> x.<warning descr="Method invocation 'trim' will produce 'NullPointerException'">trim</warning>(), Collectors.toList()));
+    Map<String, List<String>> map3 = list.stream()
+      .filter(Objects::isNull)
+      .collect(Collectors.groupingBy(<warning descr="Passing 'null' argument to parameter annotated as @NotNull">null</warning>, <warning descr="Passing 'null' argument to parameter annotated as @NotNull">null</warning>, <warning descr="Passing 'null' argument to parameter annotated as @NotNull">null</warning>));
+    Map<String, List<String>> map4 = list.stream()
+      .filter(Objects::isNull)
+      .collect(Collectors.groupingBy(x -> <warning descr="Function may return null, but it's not allowed here">null</warning>));
+    Map<Boolean, List<String>> map5 = list.stream()
+      .filter(x -> !x.isEmpty())
+      .collect(Collectors.partitioningBy(x -> <warning descr="Result of 'x.isEmpty()' is always 'false'">x.isEmpty()</warning>));
+    Map<Boolean, List<String>> map6 = list.stream()
+      .filter(x -> x.isEmpty())
+      .collect(Collectors.partitioningBy(x -> <warning descr="Result of 'x.isEmpty()' is always 'true'">x.isEmpty()</warning>, Collectors.toList()));
+    Map<Boolean, List<String>> map7 = list.stream()
+      .filter(x -> !x.isEmpty())
+      .collect(Collectors.partitioningBy(x -> <warning descr="Result of 'list.isEmpty()' is always 'false'">list.isEmpty()</warning>, Collectors.toList()));
+  }
+}

java/java-tests/testData/inspection/optionalGet/StreamGroupingBy.java

@@ -0,0 +1,40 @@
+import java.util.*;
+import java.util.stream.*;
+
+class TestClassContainingOptional {
+
+  void test() {
+    List<TestClassContainingOptional> list = new ArrayList<>();
+    list.add(new TestClassContainingOptional("name1", "optional"));
+    list.add(new TestClassContainingOptional("name2"));
+    System.out.println(list.stream()
+                         .filter(e -> e.getOptionalString().isPresent())
+                         .collect(Collectors.groupingBy(e -> e.getOptionalString().get()))
+    );
+    System.out.println(list.stream()
+                         .filter(e -> e.getOptionalString().isPresent())
+                         .collect(Collectors.groupingBy(e -> e.getOptionalString().get(), Collectors.toSet()))
+    );
+    TreeMap<Object, Set<Object>> collect = list.stream()
+      .filter(e -> e.getOptionalString().isPresent())
+      .collect(Collectors.groupingBy(e -> e.getOptionalString().get(), TreeMap::new, Collectors.toSet()));
+    System.out.println(collect);
+  }
+
+  String name;
+  String optionalString = null;
+
+
+  public TestClassContainingOptional(String name) {
+    this.name = name;
+  }
+
+  public TestClassContainingOptional(String name, String optionalString) {
+    this.name = name;
+    this.optionalString = optionalString;
+  }
+
+  public Optional<String> getOptionalString() {
+    return Optional.ofNullable(optionalString);
+  }
+}

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

@@ -218,6 +218,7 @@ public class DataFlowInspection8Test extends DataFlowInspectionTestCase {
   public void testStreamCustomSumMethod() { doTest(); }
   public void testStreamReduceLogicalAnd() { doTest(); }
   public void testStreamSingleElementReduce() { doTest(); }
+  public void testStreamGroupingBy() { doTest(); }
   public void testRequireNonNullMethodRef() {
     doTestWith((dfa, __) -> dfa.SUGGEST_NULLABLE_ANNOTATIONS = true);
   }

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

@@ -18,6 +18,7 @@ public class OptionalGetWithoutIsPresentInspectionTest extends LightJavaCodeInsi
   public void testOptionalGet() { doTest(); }
   public void testOptionalGetInlineLambda() { doTest(); }
   public void testOptionalGetMethodReference() { doTest(); }
+  public void testStreamGroupingBy() { doTest(); }
 
   @NotNull
   @Override