[java-inspections] IDEA-194953 Optional::get without isPresent inspection doesn't work in Stream::map with method reference

GitOrigin-RevId: bb857fbf79fd224406b37ad2a40dd33445575514

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

@@ -3,6 +3,9 @@ package com.intellij.codeInspection.dataFlow;
 
 import com.intellij.codeInspection.dataFlow.interpreter.RunnerResult;
 import com.intellij.codeInspection.dataFlow.java.JavaDfaListener;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaDfaAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaExpressionAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaMethodReferenceArgumentAnchor;
 import com.intellij.codeInspection.dataFlow.jvm.SpecialField;
 import com.intellij.codeInspection.dataFlow.jvm.descriptors.AssertionDisabledDescriptor;
 import com.intellij.codeInspection.dataFlow.jvm.problems.ContractFailureProblem;
@@ -89,8 +92,8 @@ public final class CommonDataflow {
    * Represents the result of dataflow applied to some code fragment (usually a method)
    */
   public static final class DataflowResult {
-    private final @NotNull Map<PsiExpression, DataflowPoint> myData = new HashMap<>();
-    private @NotNull Map<PsiExpression, DataflowPoint> myDataAssertionsDisabled = myData;
+    private final @NotNull Map<JavaDfaAnchor, DataflowPoint> myData = new HashMap<>();
+    private @NotNull Map<JavaDfaAnchor, DataflowPoint> myDataAssertionsDisabled = myData;
     private final RunnerResult myResult;
 
     public DataflowResult(RunnerResult result) {
@@ -100,45 +103,46 @@ public final class CommonDataflow {
     @NotNull
     DataflowResult copy() {
       DataflowResult copy = new DataflowResult(myResult);
-      myData.forEach((expression, point) -> copy.myData.put(expression, new DataflowPoint(point)));
+      myData.forEach((anchor, point) -> copy.myData.put(anchor, new DataflowPoint(point)));
       return copy;
     }
 
-    void add(PsiExpression expression, DfaMemoryState memState, DfaValue value) {
+    void add(JavaDfaAnchor anchor, DfaMemoryState memState, DfaValue value) {
       DfaVariableValue assertionDisabled = AssertionDisabledDescriptor.getAssertionsDisabledVar(value.getFactory());
       if (assertionDisabled == null) {
         assert myData == myDataAssertionsDisabled;
-        updateDataPoint(myData, expression, memState, value);
+        updateDataPoint(myData, anchor, memState, value);
       } else {
         DfType type = memState.getDfType(assertionDisabled);
         if (type == DfTypes.TRUE || type == DfTypes.FALSE) {
           if (myData == myDataAssertionsDisabled) {
             myDataAssertionsDisabled = new HashMap<>(myData);
           }
-          updateDataPoint(type == DfTypes.TRUE ? myDataAssertionsDisabled : myData, expression, memState, value);
+          updateDataPoint(type == DfTypes.TRUE ? myDataAssertionsDisabled : myData, anchor, memState, value);
         } else {
-          updateDataPoint(myData, expression, memState, value);
+          updateDataPoint(myData, anchor, memState, value);
           if (myData != myDataAssertionsDisabled) {
-            updateDataPoint(myDataAssertionsDisabled, expression, memState, value);
+            updateDataPoint(myDataAssertionsDisabled, anchor, memState, value);
           }
         }
       }
     }
 
-    private void updateDataPoint(Map<PsiExpression, DataflowPoint> data,
-                                 PsiExpression expression,
+    private void updateDataPoint(Map<JavaDfaAnchor, DataflowPoint> data,
+                                 JavaDfaAnchor anchor,
                                  DfaMemoryState memState,
                                  DfaValue value) {
-      DataflowPoint point = data.computeIfAbsent(expression, e -> new DataflowPoint());
+      DataflowPoint point = data.computeIfAbsent(anchor, e -> new DataflowPoint());
       if (DfaTypeValue.isContractFail(value)) {
         point.myMayFailByContract = true;
         return;
       }
-      if (point.myDfType != DfType.TOP) {
+      if (point.myDfType != DfType.TOP && anchor instanceof JavaExpressionAnchor) {
+        PsiExpression expression = ((JavaExpressionAnchor)anchor).getExpression();
         PsiElement parent = PsiUtil.skipParenthesizedExprUp(expression.getParent());
         if (parent instanceof PsiConditionalExpression &&
             !PsiTreeUtil.isAncestor(((PsiConditionalExpression)parent).getCondition(), expression, false)) {
-          add((PsiExpression)parent, memState, value);
+          add(new JavaExpressionAnchor((PsiExpression)parent), memState, value);
         }
       }
       point.addFacts(memState, value);
@@ -154,11 +158,17 @@ public final class CommonDataflow {
      * If false is returned, it's possible that the expression exists in unreachable branch or this expression is not tracked due to
      * the dataflow implementation details.
      */
-    public boolean expressionWasAnalyzed(PsiExpression expression) {
+    @Contract("null -> false")
+    public boolean expressionWasAnalyzed(@Nullable PsiExpression expression) {
+      if (expression == null) return false;
       if (expression instanceof PsiParenthesizedExpression) {
         throw new IllegalArgumentException("Should not pass parenthesized expression");
       }
-      return myData.containsKey(expression);
+      return myData.containsKey(new JavaExpressionAnchor(expression));
+    }
+
+    public boolean anchorWasAnalyzed(@NotNull JavaDfaAnchor anchor) {
+      return myData.containsKey(anchor);
     }
 
     /**
@@ -168,8 +178,10 @@ public final class CommonDataflow {
      * @param call call to check
      * @return true if it cannot fail by contract; false if unknown or can fail
      */
-    public boolean cannotFailByContract(PsiCallExpression call) {
-      DataflowPoint point = myData.get(call);
+    @Contract("null -> false")
+    public boolean cannotFailByContract(@Nullable PsiCallExpression call) {
+      if (call == null) return false;
+      DataflowPoint point = myData.get(new JavaExpressionAnchor(call));
       return point != null && !point.myMayFailByContract;
     }
 
@@ -182,7 +194,8 @@ public final class CommonDataflow {
      */
     @NotNull
     public Set<Object> getExpressionValues(@Nullable PsiExpression expression) {
-      DataflowPoint point = myData.get(expression);
+      if (expression == null) return Collections.emptySet();
+      DataflowPoint point = myData.get(new JavaExpressionAnchor(expression));
       if (point == null) return Collections.emptySet();
       Set<Object> values = point.myPossibleValues;
       return values == null ? Collections.emptySet() : Collections.unmodifiableSet(values);
@@ -196,7 +209,14 @@ public final class CommonDataflow {
      */
     @NotNull
     public DfType getDfType(PsiExpression expression) {
-      DataflowPoint point = myData.get(expression);
+      if (expression == null) return DfType.TOP;
+      DataflowPoint point = myData.get(new JavaExpressionAnchor(expression));
+      return point == null ? DfType.TOP : point.myDfType;
+    }
+
+    @NotNull
+    public DfType getDfType(@NotNull JavaDfaAnchor anchor) {
+      DataflowPoint point = myData.get(anchor);
       return point == null ? DfType.TOP : point.myDfType;
     }
 
@@ -208,7 +228,8 @@ public final class CommonDataflow {
      */
     @NotNull
     public DfType getDfTypeNoAssertions(PsiExpression expression) {
-      DataflowPoint point = myDataAssertionsDisabled.get(expression);
+      if (expression == null) return DfType.TOP;
+      DataflowPoint point = myDataAssertionsDisabled.get(new JavaExpressionAnchor(expression));
       return point == null ? DfType.TOP : point.myDfType;
     }
   }
@@ -346,7 +367,14 @@ public final class CommonDataflow {
     public void beforeExpressionPush(@NotNull DfaValue value,
                                      @NotNull PsiExpression expression,
                                      @NotNull DfaMemoryState state) {
-      myResult.add(expression, state, value);
+      myResult.add(new JavaExpressionAnchor(expression), state, value);
+    }
+
+    @Override
+    public void beforeMethodReferenceArgumentPush(@NotNull DfaValue value,
+                                                  @NotNull PsiMethodReferenceExpression expression,
+                                                  @NotNull DfaMemoryState state) {
+      myResult.add(new JavaMethodReferenceArgumentAnchor(expression), state, value);
     }
 
     @Override
@@ -355,7 +383,8 @@ public final class CommonDataflow {
                             @NotNull ThreeState failed,
                             @NotNull DfaMemoryState state) {
       if (problem instanceof ContractFailureProblem && failed != ThreeState.NO) {
-        myResult.add(((ContractFailureProblem)problem).getAnchor(), state, value.getFactory().fromDfType(DfType.FAIL));
+        myResult.add(new JavaExpressionAnchor(((ContractFailureProblem)problem).getAnchor()), state,
+                     value.getFactory().fromDfType(DfType.FAIL));
       }
     }
   }

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

@@ -5,6 +5,7 @@ import com.intellij.codeInsight.Nullability;
 import com.intellij.codeInspection.dataFlow.DfaPsiUtil;
 import com.intellij.codeInspection.dataFlow.NullabilityProblemKind;
 import com.intellij.codeInspection.dataFlow.java.anchor.JavaExpressionAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaMethodReferenceArgumentAnchor;
 import com.intellij.codeInspection.dataFlow.java.anchor.JavaMethodReferenceReturnAnchor;
 import com.intellij.codeInspection.dataFlow.java.inliner.CallInliner;
 import com.intellij.codeInspection.dataFlow.java.inst.*;
@@ -737,6 +738,9 @@ public class CFGBuilder {
       JavaResolveResult resolveResult = methodRef.advancedResolve(false);
       PsiMethod method = ObjectUtils.tryCast(resolveResult.getElement(), PsiMethod.class);
       if (method != null && !method.isVarArgs()) {
+        if (argCount == 1) {
+          add(new ResultOfInstruction(new JavaMethodReferenceArgumentAnchor(methodRef)));
+        }
         if (processKnownMethodReference(argCount, methodRef, method)) return this;
         int expectedArgCount = method.getParameterList().getParametersCount();
         boolean pushQualifier = true;

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

@@ -1,10 +1,7 @@
 // Copyright 2000-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.codeInspection.dataFlow.java;
 
-import com.intellij.codeInspection.dataFlow.java.anchor.JavaDfaAnchor;
-import com.intellij.codeInspection.dataFlow.java.anchor.JavaEndOfInstanceInitializerAnchor;
-import com.intellij.codeInspection.dataFlow.java.anchor.JavaExpressionAnchor;
-import com.intellij.codeInspection.dataFlow.java.anchor.JavaMethodReferenceReturnAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.*;
 import com.intellij.codeInspection.dataFlow.lang.DfaAnchor;
 import com.intellij.codeInspection.dataFlow.lang.DfaListener;
 import com.intellij.codeInspection.dataFlow.memory.DfaMemoryState;
@@ -37,6 +34,9 @@ public interface JavaDfaListener extends DfaListener {
       PsiExpression psiAnchor = ((JavaExpressionAnchor)anchor).getExpression();
       callBeforeExpressionPush(value, psiAnchor, state, psiAnchor);
     }
+    if (anchor instanceof JavaMethodReferenceArgumentAnchor) {
+      beforeMethodReferenceArgumentPush(value, ((JavaMethodReferenceArgumentAnchor)anchor).getMethodReference(), state);
+    }
   }
 
   /**
@@ -90,14 +90,26 @@ public interface JavaDfaListener extends DfaListener {
 
   }
 
+  /**
+   * Called before implicit sole argument of method reference is pushed to the stack
+   * @param value value that is about to be pushed
+   * @param expression corresponding method reference
+   * @param state memory state
+   */
+  default void beforeMethodReferenceArgumentPush(@NotNull DfaValue value,
+                                                 @NotNull PsiMethodReferenceExpression expression,
+                                                 @NotNull DfaMemoryState state) {
+
+  }
+
   /**
    * Called before returning the value from specific computation scope (method, lambda, method reference).
    * Can be called many times for the same scope.
-   * 
-   * @param value value to be returned
-   * @param expression expression that resulted in a given value (can be null) 
-   * @param context context ({@link PsiMethod}, or {@link PsiFunctionalExpression})
-   * @param state memory state
+   *
+   * @param value      value to be returned
+   * @param expression expression that resulted in a given value (can be null)
+   * @param context    context ({@link PsiMethod}, or {@link PsiFunctionalExpression})
+   * @param state      memory state
    */
   default void beforeValueReturn(@NotNull DfaValue value,
                                  @Nullable PsiExpression expression,

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/java/anchor/JavaMethodReferenceArgumentAnchor.java

@@ -0,0 +1,41 @@
+// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.codeInspection.dataFlow.java.anchor;
+
+import com.intellij.psi.PsiMethodReferenceExpression;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Objects;
+
+/**
+ * Implicit sole argument of the method reference pushed to the stack.
+ * Anchored only in some cases when this method reference is handled by inliner
+ */
+public class JavaMethodReferenceArgumentAnchor extends JavaDfaAnchor {
+  private final @NotNull PsiMethodReferenceExpression myMethodRef;
+
+  public JavaMethodReferenceArgumentAnchor(@NotNull PsiMethodReferenceExpression ref) {
+    myMethodRef = ref;
+  }
+
+  public @NotNull PsiMethodReferenceExpression getMethodReference() {
+    return myMethodRef;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    JavaMethodReferenceArgumentAnchor anchor = (JavaMethodReferenceArgumentAnchor)o;
+    return Objects.equals(myMethodRef, anchor.myMethodRef);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(myMethodRef);
+  }
+
+  @Override
+  public String toString() {
+    return "argument of " + myMethodRef.getText();
+  }
+}

java/java-impl/src/com/intellij/codeInspection/java18api/OptionalGetWithoutIsPresentInspection.java

@@ -5,6 +5,9 @@ import com.intellij.codeInsight.PsiEquivalenceUtil;
 import com.intellij.codeInspection.*;
 import com.intellij.codeInspection.dataFlow.CommonDataflow;
 import com.intellij.codeInspection.dataFlow.DfaNullability;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaDfaAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaExpressionAnchor;
+import com.intellij.codeInspection.dataFlow.java.anchor.JavaMethodReferenceArgumentAnchor;
 import com.intellij.codeInspection.dataFlow.jvm.SpecialField;
 import com.intellij.codeInspection.dataFlow.types.DfReferenceType;
 import com.intellij.codeInspection.dataFlow.types.DfType;
@@ -17,6 +20,7 @@ import com.intellij.psi.codeStyle.VariableKind;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.intellij.util.ObjectUtils;
 import com.siyeh.ig.psiutils.CommentTracker;
 import com.siyeh.ig.psiutils.ExpressionUtils;
 import com.siyeh.ig.psiutils.VariableNameGenerator;
@@ -39,19 +43,31 @@ public class OptionalGetWithoutIsPresentInspection extends AbstractBaseJavaLocal
         if (qualifier == null) return;
         PsiClass optionalClass = PsiUtil.resolveClassInClassTypeOnly(qualifier.getType());
         if (optionalClass == null) return;
-        CommonDataflow.DataflowResult result = CommonDataflow.getDataflowResult(qualifier);
-        if (result == null || !result.expressionWasAnalyzed(qualifier)) return;
-        DfType dfType = SpecialField.OPTIONAL_VALUE.getFromQualifier(result.getDfType(qualifier));
-        if (dfType != DfType.TOP && !(dfType instanceof DfReferenceType)) return;
-        DfaNullability nullability = DfaNullability.fromDfType(dfType);
-        if ((nullability == DfaNullability.UNKNOWN || nullability == DfaNullability.NULLABLE) &&
-            !isPresentCallWithSameQualifierExists(qualifier)) {
+        JavaExpressionAnchor anchor = new JavaExpressionAnchor(qualifier);
+        if (isOptionalProblem(qualifier, anchor) && !isPresentCallWithSameQualifierExists(qualifier)) {
           holder.registerProblem(nameElement,
                                  JavaBundle.message("inspection.optional.get.without.is.present.message", optionalClass.getName()),
                                  tryCreateFix(call));
         }
       }
 
+      private boolean isOptionalProblem(@NotNull PsiExpression context, @NotNull JavaDfaAnchor anchor) {
+        CommonDataflow.DataflowResult result = CommonDataflow.getDataflowResult(context);
+        if (result == null || !result.anchorWasAnalyzed(anchor)) return false;
+        DfType dfType = SpecialField.OPTIONAL_VALUE.getFromQualifier(result.getDfType(anchor));
+        if (dfType != DfType.TOP && !(dfType instanceof DfReferenceType)) return false;
+        DfaNullability nullability = DfaNullability.fromDfType(dfType);
+        return nullability == DfaNullability.UNKNOWN || nullability == DfaNullability.NULLABLE;
+      }
+
+      @Override
+      public void visitMethodReferenceExpression(PsiMethodReferenceExpression methodRef) {
+        if (!OptionalUtil.OPTIONAL_GET.methodReferenceMatches(methodRef)) return;
+        if (isOptionalProblem(methodRef, new JavaMethodReferenceArgumentAnchor(methodRef))) {
+          holder.registerProblem(methodRef, JavaBundle.message("inspection.optional.get.without.is.present.method.reference.message"));
+        }
+      }
+
       public boolean isPresentCallWithSameQualifierExists(PsiExpression qualifier) {
         // Conservatively skip the results of method calls if there's an isPresent() call with the same qualifier in the method
         if (qualifier instanceof PsiMethodCallExpression) {

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

@@ -0,0 +1,15 @@
+import java.util.stream.*;
+import java.util.*;
+
+public class OptionalGetMethodReference {
+  public static void main(String[] args) {
+    Stream.of(Optional.ofNullable(Math.random() > 0.5 ? 1 : null))
+      .map(<warning descr="'Optional::get' without 'isPresent()' check">Optional::get</warning>)
+      .forEach(System.out::println);
+
+    Stream.of(Optional.ofNullable(Math.random() > 0.5 ? 1 : null))
+      .filter(Optional::isPresent)
+      .map(Optional::get)
+      .forEach(System.out::println);
+  }
+}

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

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

java/openapi/resources/messages/JavaBundle.properties

@@ -551,6 +551,7 @@ inspection.nullable.problems.not.annotated.getters.for.annotated.fields=Report n
 inspection.nullable.problems.notnull.overrides.option=Report @NotNull &parameters overriding non-annotated
 inspection.nullable.problems.notnull.parameters.with.null.literal.option=Report @NotNull parameters with null-literal argument usages
 inspection.optional.get.without.is.present.message=<code>{0}.#ref()</code> without ''isPresent()'' check
+inspection.optional.get.without.is.present.method.reference.message=<code>#ref</code> without 'isPresent()' check
 inspection.overflowing.loop.index.inspection.description=Loop executes zero or billions of times
 inspection.overflowing.loop.index.inspection.name=Loop executes zero or billions of times
 inspection.overwritten.key.map.message=Duplicate Map key