[java-inspections] IDEA-379917 .filter(Objects::nonNull) should change type of Stream to nonNull

GitOrigin-RevId: 6f4d415288b813161c7627cb8839e274e189a1b8

java/java-analysis-impl/resources/META-INF/JavaAnalysisPlugin.xml

@@ -368,5 +368,6 @@
 
     <registryKey key="batch.inspections.ignored.jvm.languages" defaultValue="Groovy,Scala"
                  description="Disable UAST processing for JVM languages; a comma separated list of UAST capable language ids."/>
+    <java.expressionTypeNullabilityPatcher implementation="com.intellij.psi.impl.StreamNullabilityPatcher"/>
   </extensions>
 </idea-plugin>

java/java-analysis-impl/src/com/intellij/psi/impl/StreamNullabilityPatcher.java

@@ -0,0 +1,34 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.psi.impl;
+
+import com.intellij.codeInsight.JavaExpressionTypeNullabilityPatcher;
+import com.intellij.codeInsight.TypeNullability;
+import com.intellij.psi.*;
+import com.siyeh.ig.callMatcher.CallMatcher;
+import com.siyeh.ig.psiutils.ExpressionUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+final class StreamNullabilityPatcher implements JavaExpressionTypeNullabilityPatcher {
+  private static final CallMatcher STREAM_FILTER_TAKE_WHILE = CallMatcher.instanceCall(
+      CommonClassNames.JAVA_UTIL_STREAM_STREAM, "filter", "takeWhile")
+    .parameterCount(1);
+
+  @Override
+  public @Nullable PsiType tryPatchType(@NotNull PsiExpression expression, @NotNull PsiType type) {
+    if (type instanceof PsiClassType classType &&
+        expression instanceof PsiMethodCallExpression call &&
+        STREAM_FILTER_TAKE_WHILE.test(call) && 
+        ExpressionUtils.isNullFilteringFunction(call.getArgumentList().getExpressions()[0])) {
+      PsiClass psiClass = classType.resolve();
+      if (psiClass != null && CommonClassNames.JAVA_UTIL_STREAM_STREAM.equals(psiClass.getQualifiedName())) {
+        PsiType[] parameters = classType.getParameters();
+        if (parameters.length == 1) {
+          return JavaPsiFacade.getElementFactory(expression.getProject())
+            .createType(psiClass, parameters[0].withNullability(TypeNullability.NOT_NULL_KNOWN));
+        }
+      }
+    }
+    return null;
+  }
+}

java/java-analysis-impl/src/com/siyeh/ig/psiutils/ExpressionUtils.java

@@ -3,9 +3,7 @@ package com.siyeh.ig.psiutils;
 
 import com.intellij.codeInsight.AnnotationUtil;
 import com.intellij.codeInsight.CodeInsightUtilCore;
-import com.intellij.codeInspection.dataFlow.ContractReturnValue;
-import com.intellij.codeInspection.dataFlow.JavaMethodContractUtil;
-import com.intellij.codeInspection.dataFlow.MutationSignature;
+import com.intellij.codeInspection.dataFlow.*;
 import com.intellij.java.codeserver.core.JavaPsiReferenceUtil;
 import com.intellij.java.codeserver.core.JavaPsiReferenceUtil.ForwardReferenceProblem;
 import com.intellij.java.syntax.parser.JavaKeywords;
@@ -30,6 +28,7 @@ import org.jetbrains.annotations.*;
 
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.IntPredicate;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -1656,4 +1655,60 @@ public final class ExpressionUtils {
     }
     return codeBlock.getStatementCount() == 1;
   }
+
+  /**
+   * @param expression expression to test
+   * @return true if this expression represents a predicate-like function (accepting one parameter and returning a boolean value),
+   * which returns false for null values (may return false for some non-null values as well). This method is pattern-based,
+   * so it may not recognize specific patterns.
+   */
+  public static boolean isNullFilteringFunction(@NotNull PsiExpression expression) {
+    expression = PsiUtil.skipParenthesizedExprDown(expression);
+    if (expression instanceof PsiLambdaExpression lambda) {
+      PsiParameter[] parameters = lambda.getParameterList().getParameters();
+      if (parameters.length != 1) return false;
+      PsiParameter parameter = parameters[0];
+      PsiExpression expr = PsiUtil.skipParenthesizedExprDown(LambdaUtil.extractSingleExpressionFromBody(lambda.getBody()));
+      while (expr instanceof PsiPolyadicExpression polyadicExpression) {
+        if (polyadicExpression.getOperationTokenType() == JavaTokenType.ANDAND) {
+          expr = polyadicExpression.getOperands()[0];
+          continue;
+        }
+        return getVariableFromNullComparison(expr, false) == parameter;
+      }
+      PsiExpression negated = BoolUtils.getNegated(expr);
+      boolean nullTrue = false;
+      if (negated != null) {
+        expr = negated;
+        nullTrue = true;
+      }
+      if (expr instanceof PsiInstanceOfExpression instanceOf && isReferenceTo(instanceOf.getOperand(), parameter)) {
+        return !nullTrue;
+      }
+      if (expr instanceof PsiMethodCallExpression call) {
+        PsiMethod method = call.resolveMethod();
+        PsiExpression[] args = call.getArgumentList().getExpressions();
+        if (args.length == 0) return false;
+        return isNullFilteringMethod(method, nullTrue, arg -> isReferenceTo(args[arg], parameter));
+      }
+    }
+    if (expression instanceof PsiMethodReferenceExpression methodRef && methodRef.resolve() instanceof PsiMethod method) {
+      return isNullFilteringMethod(method, false, arg -> arg == 0);
+    }
+    return false;
+  }
+
+  private static boolean isNullFilteringMethod(@Nullable PsiMethod method, boolean nullTrue, @NotNull IntPredicate isWantedArg) {
+    if (method == null || !JavaMethodContractUtil.isPure(method)) return false;
+    List<? extends MethodContract> contracts = JavaMethodContractUtil.getMethodCallContracts(method, null);
+    for (MethodContract contract : contracts) {
+      if (!(ContainerUtil.getOnlyItem(contract.getConditions()) instanceof ContractValue.Condition condition)) return false;
+      int argEqNull = condition.getArgumentComparedTo(ContractValue.nullValue(), true).orElse(-1);
+      if (argEqNull != -1 && isWantedArg.test(argEqNull)) {
+        ContractReturnValue value = contract.getReturnValue();
+        return value.equals(ContractReturnValue.returnBoolean(nullTrue));
+      }
+    }
+    return false;
+  }
 }

java/java-psi-api/src/com/intellij/codeInsight/JavaExpressionTypeNullabilityPatcher.java

@@ -0,0 +1,48 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight;
+
+import com.intellij.openapi.extensions.ExtensionPointName;
+import com.intellij.psi.PsiExpression;
+import com.intellij.psi.PsiType;
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * An extension point which allows to patch the type nullability of expressions.
+ */
+@ApiStatus.Experimental
+public interface JavaExpressionTypeNullabilityPatcher {
+  ExtensionPointName<JavaExpressionTypeNullabilityPatcher> EP_NAME = ExtensionPointName.create("com.intellij.java.expressionTypeNullabilityPatcher");
+  
+  /**
+   * Patches the type nullability of the given expression using the extra-linguistic knowledge
+   * (e.g., call to a known library method). May patch not only the top-level nullability but
+   * also nested nullability as well. Implement with care, as this is computed when determining 
+   * the expression type, so it's easy to get into infinite recursion if other non-trivial operations
+   * are performed on the expression.
+   * 
+   * @param expression expression
+   * @param type computed type with inherent nullability
+   * @return patched type, or null if no patching is supported for a given expression. In this case,
+   * the next patcher will be tried, if available.
+   */
+  @Nullable PsiType tryPatchType(@NotNull PsiExpression expression, @NotNull PsiType type);
+
+  /**
+   * Patches the type nullability of the given expression using all registered patchers.
+   * 
+   * @param expression expression whose type should be patched
+   * @param type computed type with inherent nullability
+   * @return the patched type, or the originally computed type if no patcher wants to patch this expression
+   */
+  static @NotNull PsiType patchTypeNullability(@NotNull PsiExpression expression, @NotNull PsiType type) {
+    for (JavaExpressionTypeNullabilityPatcher patcher : EP_NAME.getExtensionList()) {
+      PsiType patchedType = patcher.tryPatchType(expression, type);
+      if (patchedType != null) {
+        return patchedType;
+      }
+    }
+    return type;
+  } 
+}

java/java-psi-api/src/com/intellij/codeInsight/NullabilitySource.java

@@ -33,6 +33,12 @@ public /* sealed */ interface NullabilitySource {
      * (e.g., primitive type, or disjunction type)
      */
     MANDATED,
+    /**
+     * Type nullability is known from a particular code shape.
+     * While it may differ from one defined by the language constructs only,
+     * it's believed to be correct and more helpful to users. 
+     */
+    KNOWN,
     /**
      * Type nullability is depicted explicitly by means of the language.
      * Currently, not possible in Java, but may be used in other languages like Kotlin.

java/java-psi-api/src/com/intellij/codeInsight/TypeNullability.java

@@ -25,6 +25,10 @@ public final class TypeNullability {
    * Mandated not-null nullability
    */
   public static final TypeNullability NOT_NULL_MANDATED = new TypeNullability(Nullability.NOT_NULL, NullabilitySource.Standard.MANDATED);
+  /**
+   * Known not-null nullability
+   */
+  public static final TypeNullability NOT_NULL_KNOWN = new TypeNullability(Nullability.NOT_NULL, NullabilitySource.Standard.KNOWN);
   /**
    * Mandated nullable nullability
    */

java/java-psi-impl/resources/intellij.java.psi.impl.xml

@@ -33,6 +33,7 @@
     <extensionPoint qualifiedName="com.intellij.javadocTagInfo" area="IDEA_PROJECT" interface="com.intellij.psi.javadoc.JavadocTagInfo" dynamic="true"/>
     <extensionPoint qualifiedName="com.intellij.importFilter" interface="com.intellij.codeInsight.ImportFilter" dynamic="true"/>
     <extensionPoint qualifiedName="com.intellij.javaMainMethodProvider" interface="com.intellij.codeInsight.runner.JavaMainMethodProvider" dynamic="true"/>
+    <extensionPoint qualifiedName="com.intellij.java.expressionTypeNullabilityPatcher" interface="com.intellij.codeInsight.JavaExpressionTypeNullabilityPatcher" dynamic="true"/>
     <extensionPoint qualifiedName="com.intellij.superMethodsSearch" interface="com.intellij.util.QueryExecutor" dynamic="true"/>
     <extensionPoint qualifiedName="com.intellij.deepestSuperMethodsSearch" interface="com.intellij.util.QueryExecutor" dynamic="true"/>
     <extensionPoint qualifiedName="com.intellij.lang.jvm.multiReleaseSupport" interface="com.intellij.psi.util.JavaMultiReleaseModuleSupport" dynamic="true"/>

java/java-psi-impl/src/com/intellij/psi/impl/source/tree/java/PsiMethodCallExpressionImpl.java

@@ -1,6 +1,7 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.psi.impl.source.tree.java;
 
+import com.intellij.codeInsight.JavaExpressionTypeNullabilityPatcher;
 import com.intellij.codeInsight.daemon.impl.analysis.JavaGenericsUtil;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.diagnostic.Logger;
@@ -169,7 +170,8 @@ public class PsiMethodCallExpressionImpl extends ExpressionPsiElement implements
         }
       }
 
-      return PsiClassImplUtil.correctType(theOnly, file.getResolveScope());
+      PsiType correctedType = PsiClassImplUtil.correctType(theOnly, file.getResolveScope());
+      return correctedType == null ? null : JavaExpressionTypeNullabilityPatcher.patchTypeNullability(call, correctedType);
     }
 
     private static @Nullable PsiType getResultType(@NotNull PsiMethodCallExpression call,

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

@@ -0,0 +1,43 @@
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+import java.util.*;
+
+@NullMarked
+public class StreamFilterPatching {
+  public static List<Integer> getIntegersMR(List<@Nullable Integer> integers) {
+    return integers.stream().filter(Objects::nonNull).toList();
+  }
+
+  public static List<Integer> getIntegersMR2(List<@Nullable Integer> integers) {
+    return <warning descr="Returning a class with nullable type parameters when a class with non-null type parameters is expected">integers.stream().filter(Objects::isNull).toList()</warning>;
+  }
+
+  public static List<Integer> getIntegersLambda(List<@Nullable Integer> integers) {
+    return integers.stream().filter(obj -> Objects.nonNull(obj)).toList();
+  }
+
+  public static List<Integer> getIntegersLambda2(List<@Nullable Integer> integers) {
+    return <warning descr="Returning a class with nullable type parameters when a class with non-null type parameters is expected">integers.stream().filter(obj -> Objects.isNull(obj)).toList()</warning>;
+  }
+
+  public static List<Integer> getIntegersLambda3(List<@Nullable Integer> integers) {
+    return integers.stream().filter(obj -> !Objects.isNull(obj)).toList();
+  }
+
+  public static List<Integer> getIntegersLambda4(List<@Nullable Integer> integers) {
+    return integers.stream().filter(obj -> obj != null).toList();
+  }
+
+  public static List<Number> getIntegersLambda5(List<@Nullable Number> integers) {
+    return integers.stream().filter(obj -> obj instanceof Integer).toList();
+  }
+
+  public static List<Number> getIntegersLambda5And(List<@Nullable Number> integers) {
+    return integers.stream().filter(obj -> obj instanceof Integer i && i == 10).toList();
+  }
+
+  public static List<Number> getIntegersInstance(List<@Nullable Number> integers) {
+    return integers.stream().filter(Integer.class::isInstance).toList();
+  }
+}

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

@@ -492,6 +492,12 @@ public class NullableStuffInspectionTest extends LightJavaCodeInsightFixtureTest
     doTest();
   }
   
+  public void testStreamFilterPatching() {
+    addJSpecifyNullMarked(myFixture);
+    setupTypeUseAnnotations("org.jspecify.annotations", myFixture);
+    doTest();
+  }
+
   public void testNullableExtendsNullable() {
     addJSpecifyNullMarked(myFixture);
     setupTypeUseAnnotations("org.jspecify.annotations", myFixture);