DFA instruction visitor refactoring wave#3

Always failing calls, constant calls moved to beforeExpressionPush handler
JavaMethodContractUtil API tweaks
DfaConstValue#isContractFail

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

@@ -93,7 +93,7 @@ public class CommonDataflow {
   private static DataflowResult runDFA(@Nullable PsiElement block) {
     if (block == null) return null;
     DataFlowRunner runner = new DataFlowRunner(false, block);
-    CommonDataflowVisitor visitor = new CommonDataflowVisitor(runner);
+    CommonDataflowVisitor visitor = new CommonDataflowVisitor();
     RunnerResult result = runner.analyzeMethodRecursively(block, visitor);
     if (result != RunnerResult.OK) return null;
     if (!(block instanceof PsiClass)) return visitor.myResult;
@@ -148,15 +148,9 @@ public class CommonDataflow {
   }
 
   private static class CommonDataflowVisitor extends StandardInstructionVisitor {
-    private DataflowResult myResult;
+    private DataflowResult myResult = new DataflowResult();
-    private final DfaConstValue myFail;
     private final List<DfaMemoryState> myEndOfInitializerStates = new ArrayList<>();
 
-    public CommonDataflowVisitor(DataFlowRunner runner) {
-      myFail = runner.getFactory().getConstFactory().getContractFail();
-      myResult = new DataflowResult();
-    }
-
     @Override
     public DfaInstructionState[] visitEndOfInitializer(EndOfInitializerInstruction instruction,
                                                        DataFlowRunner runner,
@@ -172,7 +166,7 @@ public class CommonDataflow {
                                      @NotNull PsiExpression expression,
                                      @Nullable TextRange range,
                                      @NotNull DfaMemoryState state) {
-      if (range == null && value != myFail) {
+      if (range == null && !DfaConstValue.isContractFail(value)) {
         // Do not track instructions which cover part of expression
         myResult.add(expression, (DfaMemoryStateImpl)state, value);
       }

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

@@ -15,7 +15,6 @@
  */
 package com.intellij.codeInspection.dataFlow;
 
-import com.intellij.codeInspection.dataFlow.instructions.MethodCallInstruction;
 import com.intellij.codeInspection.dataFlow.rangeSet.LongRangeSet;
 import com.intellij.codeInspection.dataFlow.value.DfaConstValue;
 import com.intellij.codeInspection.dataFlow.value.DfaValue;
@@ -76,8 +75,7 @@ class CustomMethodHandlers {
     .register(DfaOptionalSupport.OPTIONAL_OF_NULLABLE,
               (args, memState, factory) -> ofNullable(args.myArguments[0], memState, factory));
 
-  public static CustomMethodHandler find(MethodCallInstruction instruction) {
+  public static CustomMethodHandler find(PsiMethod method) {
-    PsiMethod method = instruction.getTargetMethod();
     CustomMethodHandler handler = null;
     if (isConstantCall(method)) {
       handler = (args, memState, factory) -> handleConstantCall(args, memState, factory, method);

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

@@ -437,17 +437,15 @@ public class DataFlowInspectionBase extends AbstractBaseJavaLocalInspectionTool
   private static void reportAlwaysFailingCalls(ProblemsHolder holder,
                                                DataFlowInstructionVisitor visitor,
                                                HashSet<PsiElement> reportedAnchors) {
-    visitor.getAlwaysFailingCalls().forEach((call, contracts) -> {
+    visitor.alwaysFailingCalls().remove(TestUtils::isExceptionExpected).forEach(call -> {
-      if (TestUtils.isExceptionExpected(call)) return;
+      if (reportedAnchors.add(call)) {
-      PsiMethod method = call.resolveMethod();
+        holder.registerProblem(getElementToHighlight(call), getContractMessage(JavaMethodContractUtil.getMethodCallContracts(call)));
-      if (method != null && reportedAnchors.add(call)) {
-        holder.registerProblem(getElementToHighlight(call), getContractMessage(contracts));
       }
     });
   }
 
   @NotNull
-  private static String getContractMessage(List<MethodContract> contracts) {
+  private static String getContractMessage(List<? extends MethodContract> contracts) {
     if (contracts.stream().allMatch(mc -> mc.getConditions().stream().allMatch(ContractValue::isBoundCheckingCondition))) {
       return InspectionsBundle.message("dataflow.message.contract.fail.index");
     }

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

@@ -21,12 +21,14 @@ import org.jetbrains.annotations.Nullable;
 import java.util.*;
 import java.util.stream.Stream;
 
+import static com.intellij.util.ObjectUtils.tryCast;
+
 final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
   private static final Logger LOG = Logger.getInstance("#com.intellij.codeInspection.dataFlow.DataFlowInstructionVisitor");
   private static final Object ANY_VALUE = ObjectUtils.sentinel("ANY_VALUE");
   private final Map<NullabilityProblemKind.NullabilityProblem<?>, StateInfo> myStateInfos = new LinkedHashMap<>();
   private final Set<Instruction> myCCEInstructions = ContainerUtil.newHashSet();
-  private final Map<MethodCallInstruction, Boolean> myFailingCalls = new HashMap<>();
+  private final Map<PsiCallExpression, Boolean> myFailingCalls = new HashMap<>();
   private final Map<PsiMethodCallExpression, ThreeState> myBooleanCalls = new HashMap<>();
   private final Map<PsiElement, ThreeState> myOfNullableCalls = new HashMap<>();
   private final Map<PsiAssignmentExpression, Pair<PsiType, PsiType>> myArrayStoreProblems = new HashMap<>();
@@ -152,9 +154,8 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     return StreamEx.ofKeys(myOutOfBoundsArrayAccesses, ThreeState.YES::equals);
   }
 
-  Map<PsiCall, List<MethodContract>> getAlwaysFailingCalls() {
+  StreamEx<PsiCallExpression> alwaysFailingCalls() {
-    return StreamEx.ofKeys(myFailingCalls, v -> v)
+    return StreamEx.ofKeys(myFailingCalls, v -> v);
-      .mapToEntry(MethodCallInstruction::getCallExpression, MethodCallInstruction::getContracts).toMap();
   }
 
   boolean isAlwaysReturnsNotNull(Instruction[] instructions) {
@@ -167,10 +168,7 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
                                       @NotNull PsiExpression expression,
                                       @Nullable TextRange range,
                                       @NotNull DfaMemoryState memState) {
-    if (expression instanceof PsiMethodCallExpression &&
+    expression.accept(new ExpressionVisitor(value, memState));
-        DfaOptionalSupport.OPTIONAL_OF_NULLABLE.test((PsiMethodCallExpression)expression)) {
-      processOfNullableResult(value, memState, ((PsiMethodCallExpression)expression).getArgumentList().getExpressions()[0]);
-    }
   }
 
   @Override
@@ -180,7 +178,7 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     if (DfaOptionalSupport.OPTIONAL_OF_NULLABLE.methodReferenceMatches(methodRef)) {
       processOfNullableResult(value, state, methodRef.getReferenceNameElement());
     }
-    PsiMethod method = ObjectUtils.tryCast(methodRef.resolve(), PsiMethod.class);
+    PsiMethod method = tryCast(methodRef.resolve(), PsiMethod.class);
     if (method != null) {
       List<StandardMethodContract> contracts = JavaMethodContractUtil.getMethodContracts(method);
       if (contracts.isEmpty() || !contracts.get(0).isTrivial()) {
@@ -190,47 +188,12 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     }
   }
 
-  private void processOfNullableResult(@NotNull DfaValue value,
+  private void processOfNullableResult(@NotNull DfaValue value, @NotNull DfaMemoryState memState, PsiElement anchor) {
-                                       @NotNull DfaMemoryState memState, PsiElement anchor) {
     Boolean fact = memState.getValueFact(value, DfaFactType.OPTIONAL_PRESENCE);
     ThreeState present = fact == null ? ThreeState.UNSURE : ThreeState.fromBoolean(fact);
     myOfNullableCalls.merge(anchor, present, ThreeState::merge);
   }
 
-  @Override
-  public DfaInstructionState[] visitMethodCall(MethodCallInstruction instruction,
-                                               DataFlowRunner runner,
-                                               DfaMemoryState memState) {
-    DfaInstructionState[] states = super.visitMethodCall(instruction, runner, memState);
-    if (hasNonTrivialFailingContracts(instruction)) {
-      DfaConstValue fail = runner.getFactory().getConstFactory().getContractFail();
-      boolean allFail = Arrays.stream(states).allMatch(s -> s.getMemoryState().peek() == fail);
-      myFailingCalls.merge(instruction, allFail, Boolean::logicalAnd);
-    }
-    handleBooleanCalls(instruction, states);
-    return states;
-  }
-
-  void handleBooleanCalls(MethodCallInstruction instruction, DfaInstructionState[] states) {
-    if (!hasNonTrivialBooleanContracts(instruction)) return;
-    PsiMethod method = instruction.getTargetMethod();
-    if (method == null || !JavaMethodContractUtil.isPure(method)) return;
-    PsiMethodCallExpression call = ObjectUtils.tryCast(instruction.getCallExpression(), PsiMethodCallExpression.class);
-    if (call == null || myBooleanCalls.get(call) == ThreeState.UNSURE) return;
-    if (ExpressionUtils.isVoidContext(call)) return;
-    for (DfaInstructionState s : states) {
-      DfaValue val = s.getMemoryState().peek();
-      ThreeState state = ThreeState.UNSURE;
-      if (val instanceof DfaConstValue) {
-        Object value = ((DfaConstValue)val).getValue();
-        if (value instanceof Boolean) {
-          state = ThreeState.fromBoolean((Boolean)value);
-        }
-      }
-      myBooleanCalls.merge(call, state, ThreeState::merge);
-    }
-  }
-
   @Override
   protected void processArrayAccess(PsiArrayAccessExpression expression, boolean alwaysOutOfBounds) {
     myOutOfBoundsArrayAccesses.merge(expression, ThreeState.fromBoolean(alwaysOutOfBounds), ThreeState::merge);
@@ -279,17 +242,19 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     return result;
   }
 
-  private static boolean hasNonTrivialFailingContracts(MethodCallInstruction instruction) {
+  private static boolean hasNonTrivialFailingContracts(PsiCallExpression call) {
-    List<MethodContract> contracts = instruction.getContracts();
+    List<? extends MethodContract> contracts = JavaMethodContractUtil.getMethodCallContracts(call);
-    return !contracts.isEmpty() && contracts.stream().anyMatch(
+    return !contracts.isEmpty() &&
-      contract -> contract.getReturnValue().isFail() && !contract.isTrivial());
+           contracts.stream().anyMatch(contract -> contract.getReturnValue().isFail() && !contract.isTrivial());
   }
 
-  private static boolean hasNonTrivialBooleanContracts(MethodCallInstruction instruction) {
+  private static boolean shouldCollectBooleanCallResult(PsiMethodCallExpression call) {
-    if (CustomMethodHandlers.find(instruction) != null) return true;
+    if (ExpressionUtils.isVoidContext(call)) return false;
-    List<MethodContract> contracts = instruction.getContracts();
+    PsiMethod method = call.resolveMethod();
-    return !contracts.isEmpty() && contracts.stream().anyMatch(
+    if (method == null || !PsiType.BOOLEAN.equals(method.getReturnType()) || !JavaMethodContractUtil.isPure(method)) return false;
-      contract -> contract.getReturnValue().isBoolean() && !contract.isTrivial());
+    List<? extends MethodContract> contracts = JavaMethodContractUtil.getMethodCallContracts(method, call);
+    return CustomMethodHandlers.find(method) != null ||
+           !contracts.isEmpty() && contracts.stream().anyMatch(contract -> contract.getReturnValue().isBoolean() && !contract.isTrivial());
   }
 
   @Override
@@ -336,4 +301,50 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     boolean normalNpe;
     boolean normalOk;
   }
+
+  private class ExpressionVisitor extends JavaElementVisitor {
+    private final DfaValue myValue;
+    private final DfaMemoryState myMemState;
+
+    public ExpressionVisitor(DfaValue value, DfaMemoryState memState) {
+      myValue = value;
+      myMemState = memState;
+    }
+
+    @Override
+    public void visitMethodCallExpression(PsiMethodCallExpression call) {
+      super.visitMethodCallExpression(call);
+      if (DfaOptionalSupport.OPTIONAL_OF_NULLABLE.test(call)) {
+        processOfNullableResult(myValue, myMemState, call.getArgumentList().getExpressions()[0]);
+      }
+      handleBooleanCalls(call);
+    }
+
+    @Override
+    public void visitCallExpression(PsiCallExpression call) {
+      super.visitCallExpression(call);
+      Boolean isFailing = myFailingCalls.get(call);
+      if (isFailing != null || hasNonTrivialFailingContracts(call)) {
+        myFailingCalls.put(call, DfaConstValue.isContractFail(myValue) && !Boolean.FALSE.equals(isFailing));
+      }
+    }
+
+    private void handleBooleanCalls(PsiMethodCallExpression call) {
+      ThreeState curState = myBooleanCalls.get(call);
+      if (curState == ThreeState.UNSURE) return;
+      ThreeState nextState = ThreeState.UNSURE;
+      if (myValue instanceof DfaConstValue) {
+        Object val = ((DfaConstValue)myValue).getValue();
+        if (val instanceof Boolean) {
+          nextState = ThreeState.fromBoolean((Boolean)val);
+          if (curState != null && curState != nextState) {
+            nextState = ThreeState.UNSURE;
+          }
+        }
+      }
+      if (curState != null || shouldCollectBooleanCallResult(call)) {
+        myBooleanCalls.put(call, nextState);
+      }
+    }
+  }
 }

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

@@ -2,13 +2,11 @@
 package com.intellij.codeInspection.dataFlow;
 
 import com.intellij.codeInsight.AnnotationUtil;
-import com.intellij.psi.PsiAnnotation;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiExpression;
-import com.intellij.psi.PsiMethod;
-import com.intellij.psi.PsiMethodCallExpression;
 import com.intellij.psi.util.CachedValueProvider;
 import com.intellij.psi.util.CachedValuesManager;
 import com.intellij.psi.util.PsiModificationTracker;
+import com.intellij.util.ObjectUtils;
 import com.siyeh.ig.psiutils.ExpressionUtils;
 import com.siyeh.ig.psiutils.MethodCallUtils;
 import org.jetbrains.annotations.Contract;
@@ -30,6 +28,18 @@ public class JavaMethodContractUtil {
    */
   public static final String ORG_JETBRAINS_ANNOTATIONS_CONTRACT = Contract.class.getName();
 
+  /**
+   * Returns a list of contracts defined for given method call (including hardcoded contracts if any)
+   *
+   * @param call a method call site.
+   * @return list of contracts (empty list if no contracts found)
+   */
+  @NotNull
+  public static List<? extends MethodContract> getMethodCallContracts(@NotNull PsiCallExpression call) {
+    PsiMethod method = call.resolveMethod();
+    return method == null ? Collections.emptyList() : getMethodCallContracts(method, call);
+  }
+
   /**
    * Returns a list of contracts defined for given method call (including hardcoded contracts if any)
    *
@@ -40,8 +50,9 @@ public class JavaMethodContractUtil {
    */
   @NotNull
   public static List<? extends MethodContract> getMethodCallContracts(@NotNull final PsiMethod method,
-                                                                      @Nullable PsiMethodCallExpression call) {
+                                                                      @Nullable PsiCallExpression call) {
-    List<MethodContract> contracts = HardcodedContracts.getHardcodedContracts(method, call);
+    List<MethodContract> contracts =
+      HardcodedContracts.getHardcodedContracts(method, ObjectUtils.tryCast(call, PsiMethodCallExpression.class));
     return !contracts.isEmpty() ? contracts : getMethodContracts(method);
   }
 
@@ -200,9 +211,7 @@ public class JavaMethodContractUtil {
   @Contract("null -> null")
   public static PsiExpression findReturnedValue(@Nullable PsiMethodCallExpression call) {
     if (call == null) return null;
-    PsiMethod method = call.resolveMethod();
+    List<? extends MethodContract> contracts = getMethodCallContracts(call);
-    if (method == null) return null;
-    List<? extends MethodContract> contracts = getMethodCallContracts(method, call);
     ContractReturnValue returnValue = getNonFailingReturnValue(contracts);
     if (returnValue == null) return null;
     if (returnValue.equals(ContractReturnValue.returnThis())) {

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

@@ -331,8 +331,9 @@ public class StandardInstructionVisitor extends InstructionVisitor {
 
   @NotNull
   private List<DfaMemoryState> handleKnownMethods(MethodCallInstruction instruction, DataFlowRunner runner, DfaMemoryState memState) {
-    if (instruction.getTargetMethod() == null) return Collections.emptyList();
+    PsiMethod method = instruction.getTargetMethod();
-    CustomMethodHandlers.CustomMethodHandler handler = CustomMethodHandlers.find(instruction);
+    if (method == null) return Collections.emptyList();
+    CustomMethodHandlers.CustomMethodHandler handler = CustomMethodHandlers.find(method);
     if (handler == null) return Collections.emptyList();
     memState = memState.createCopy();
     DfaCallArguments callArguments = popCall(instruction, runner, memState, false);
@@ -784,8 +785,7 @@ public class StandardInstructionVisitor extends InstructionVisitor {
     }
 
     if (dfaLeft instanceof DfaConstValue && dfaRight instanceof DfaConstValue ||
-        dfaLeft == runner.getFactory().getConstFactory().getContractFail() ||
+        DfaConstValue.isContractFail(dfaLeft) || DfaConstValue.isContractFail(dfaRight)) {
-        dfaRight == runner.getFactory().getConstFactory().getContractFail()) {
       boolean negated = (relationType == RelationType.NE) ^ (DfaMemoryStateImpl.isNaN(dfaLeft) || DfaMemoryStateImpl.isNaN(dfaRight));
       boolean result = dfaLeft == dfaRight ^ negated;
       return makeBooleanResultArray(instruction, runner, memState, result);

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/value/DfaConstValue.java

@@ -24,6 +24,7 @@ import com.intellij.psi.util.TypeConversionUtil;
 import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.ig.psiutils.ExpressionUtils;
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -190,4 +191,15 @@ public class DfaConstValue extends DfaValue {
     if (this == myFactory.getConstFactory().getFalse()) return myFactory.getConstFactory() .getTrue();
     return DfaUnknownValue.getInstance();
   }
+
+  /**
+   * Checks whether given value is a special value representing method failure, according to its contract
+   *
+   * @param value value to check
+   * @return true if specified value represents method failure
+   */
+  @Contract("null -> false")
+  public static boolean isContractFail(DfaValue value) {
+    return value instanceof DfaConstValue && ((DfaConstValue)value).getValue() == ourThrowable;
+  }
 }

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/bugs/SuspiciousComparatorCompareInspection.java

@@ -100,11 +100,8 @@ public class SuspiciousComparatorCompareInspection extends BaseInspection {
       }
       PsiMethodCallExpression soleCall = ObjectUtils.tryCast(LambdaUtil.extractSingleExpressionFromBody(body), PsiMethodCallExpression.class);
       if (soleCall != null) {
-        PsiMethod method = soleCall.resolveMethod();
+        MethodContract contract = ContainerUtil.getOnlyItem(JavaMethodContractUtil.getMethodCallContracts(soleCall));
-        if (method != null) {
+        if (contract != null && contract.isTrivial() && contract.getReturnValue().isFail()) return;
-          MethodContract contract = ContainerUtil.getOnlyItem(JavaMethodContractUtil.getMethodCallContracts(method, soleCall));
-          if (contract != null && contract.isTrivial() && contract.getReturnValue().isFail()) return;
-        }
       }
       PsiParameter[] parameters = parameterList.getParameters();
       checkParameterList(parameters, body);