DFA instruction visitor refactoring wave#2

Method reference results processing polished

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

@@ -168,7 +168,7 @@ public class CommonDataflow {
     }
 
     @Override
-    public void beforeExpressionPush(@NotNull DfaValue value,
+    protected void beforeExpressionPush(@NotNull DfaValue value,
                                      @NotNull PsiExpression expression,
                                      @Nullable TextRange range,
                                      @NotNull DfaMemoryState state) {

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

@@ -163,30 +163,38 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
   }
 
   @Override
-  public void beforeExpressionPush(@NotNull DfaValue value,
+  protected void beforeExpressionPush(@NotNull DfaValue value,
-                                   @NotNull PsiExpression expression,
+                                      @NotNull PsiExpression expression,
-                                   @Nullable TextRange range,
+                                      @Nullable TextRange range,
-                                   @NotNull DfaMemoryState memState) {
+                                      @NotNull DfaMemoryState memState) {
-    PsiElement anchor = extractOptionalOfNullableAnchor(expression);
+    if (expression instanceof PsiMethodCallExpression &&
-    if (anchor != null) {
+        DfaOptionalSupport.OPTIONAL_OF_NULLABLE.test((PsiMethodCallExpression)expression)) {
-      Boolean fact = memState.getValueFact(value, DfaFactType.OPTIONAL_PRESENCE);
+      processOfNullableResult(value, memState, ((PsiMethodCallExpression)expression).getArgumentList().getExpressions()[0]);
-      ThreeState present = fact == null ? ThreeState.UNSURE : ThreeState.fromBoolean(fact);
-      myOfNullableCalls.merge(anchor, present, ThreeState::merge);
     }
   }
 
-  private static PsiElement extractOptionalOfNullableAnchor(PsiExpression expression) {
+  @Override
-    if (expression instanceof PsiMethodCallExpression &&
+  protected void beforeMethodReferenceResultPush(@NotNull DfaValue value,
-        DfaOptionalSupport.OPTIONAL_OF_NULLABLE.test((PsiMethodCallExpression)expression)) {
+                                                 @NotNull PsiMethodReferenceExpression methodRef,
-      return ((PsiMethodCallExpression)expression).getArgumentList().getExpressions()[0];
+                                                 @NotNull DfaMemoryState state) {
+    if (DfaOptionalSupport.OPTIONAL_OF_NULLABLE.methodReferenceMatches(methodRef)) {
+      processOfNullableResult(value, state, methodRef.getReferenceNameElement());
     }
-    if (expression instanceof PsiMethodReferenceExpression) {
+    PsiMethod method = ObjectUtils.tryCast(methodRef.resolve(), PsiMethod.class);
-      PsiMethodReferenceExpression methodRef = (PsiMethodReferenceExpression)expression;
+    if (method != null) {
-      if (DfaOptionalSupport.OPTIONAL_OF_NULLABLE.methodReferenceMatches(methodRef)) {
+      List<StandardMethodContract> contracts = JavaMethodContractUtil.getMethodContracts(method);
-        return methodRef.getReferenceNameElement();
+      if (contracts.isEmpty() || !contracts.get(0).isTrivial()) {
+        // Do not track if method reference may have different results
+        myMethodReferenceResults.merge(methodRef, value, (a, b) -> a == b ? a : DfaUnknownValue.getInstance());
       }
     }
-    return null;
+  }
+
+  private void processOfNullableResult(@NotNull DfaValue value,
+                                       @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
@@ -235,16 +243,6 @@ final class DataFlowInstructionVisitor extends StandardInstructionVisitor {
     }
   }
 
-  @Override
-  protected void processMethodReferenceResult(PsiMethodReferenceExpression methodRef,
-                                              List<? extends MethodContract> contracts,
-                                              DfaValue res) {
-    if(contracts.isEmpty() || !contracts.get(0).isTrivial()) {
-      // Do not track if method reference may have different results
-      myMethodReferenceResults.merge(methodRef, res, (a, b) -> a == b ? a : DfaUnknownValue.getInstance());
-    }
-  }
-
   @Override
   public DfaInstructionState[] visitPush(PushInstruction instruction, DataFlowRunner runner, DfaMemoryState memState) {
     PsiExpression place = instruction.getExpression();

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

@@ -21,6 +21,7 @@ import com.intellij.codeInspection.dataFlow.value.*;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiArrayAccessExpression;
 import com.intellij.psi.PsiExpression;
+import com.intellij.psi.PsiMethodReferenceExpression;
 import com.intellij.psi.PsiType;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.ObjectUtils;
@@ -34,10 +35,16 @@ import java.util.ArrayList;
  */
 public abstract class InstructionVisitor {
 
-  public void beforeExpressionPush(@NotNull DfaValue value,
+  protected void beforeExpressionPush(@NotNull DfaValue value,
-                                   @NotNull PsiExpression expression,
+                                      @NotNull PsiExpression expression,
-                                   @Nullable TextRange range,
+                                      @Nullable TextRange range,
-                                   @NotNull DfaMemoryState state) {
+                                      @NotNull DfaMemoryState state) {
+
+  }
+
+  protected void beforeMethodReferenceResultPush(@NotNull DfaValue value,
+                                                 @NotNull PsiMethodReferenceExpression methodRef,
+                                                 @NotNull DfaMemoryState state) {
 
   }
 
@@ -45,8 +52,17 @@ public abstract class InstructionVisitor {
                             @NotNull ExpressionPushingInstruction instruction,
                             @NotNull DfaMemoryState state) {
     PsiExpression anchor = instruction.getExpression();
-    if (anchor != null && !(instruction instanceof PushInstruction && ((PushInstruction)instruction).isReferenceWrite())) {
+    if (anchor != null
-      beforeExpressionPush(value, anchor, instruction.getExpressionRange(), state);
+        && !(instruction instanceof MethodCallInstruction &&
+             (((MethodCallInstruction)instruction).getMethodType() == MethodCallInstruction.MethodType.BOXING ||
+              ((MethodCallInstruction)instruction).getMethodType() == MethodCallInstruction.MethodType.UNBOXING))
+        && !(instruction instanceof PushInstruction && ((PushInstruction)instruction).isReferenceWrite())) {
+      if (anchor instanceof PsiMethodReferenceExpression && !(instruction instanceof PushInstruction)) {
+        beforeMethodReferenceResultPush(value, (PsiMethodReferenceExpression)anchor, state);
+      }
+      else {
+        beforeExpressionPush(value, anchor, instruction.getExpressionRange(), state);
+      }
     }
     state.push(value);
   }

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

@@ -32,12 +32,10 @@ import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.ig.psiutils.MethodUtils;
 import com.siyeh.ig.psiutils.TypeUtils;
 import gnu.trove.THashSet;
-import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
-import java.util.stream.Stream;
 
 /**
  * @author peter
@@ -234,8 +232,13 @@ public class StandardInstructionVisitor extends InstructionVisitor {
     if (contracts.isEmpty()) return;
     PsiType returnType = substitutor.substitute(method.getReturnType());
     DfaValue defaultResult = runner.getFactory().createTypeValue(returnType, DfaPsiUtil.getElementNullability(returnType, method));
-    Stream<DfaValue> returnValues = possibleReturnValues(callArguments, state, contracts, runner.getFactory(), defaultResult, methodRef);
+    Set<DfaCallState> currentStates = Collections.singleton(new DfaCallState(state.createClosureState(), callArguments));
-    returnValues.forEach(res -> processMethodReferenceResult(methodRef, contracts, res));
+    for (MethodContract contract : contracts) {
+      currentStates = addContractResults(contract, currentStates, runner.getFactory(), new HashSet<>(), defaultResult, methodRef);
+    }
+    for (DfaCallState currentState: currentStates) {
+      pushExpressionResult(defaultResult, () -> methodRef, currentState.myMemoryState);
+    }
   }
 
   @NotNull
@@ -268,26 +271,6 @@ public class StandardInstructionVisitor extends InstructionVisitor {
     return new DfaCallArguments(qualifier, arguments, JavaMethodContractUtil.isPure(method));
   }
 
-  private Stream<DfaValue> possibleReturnValues(DfaCallArguments callArguments,
-                                                DfaMemoryState state,
-                                                List<? extends MethodContract> contracts,
-                                                DfaValueFactory factory,
-                                                DfaValue defaultResult,
-                                                PsiMethodReferenceExpression methodRef) {
-    Set<DfaCallState> currentStates = Collections.singleton(new DfaCallState(state.createClosureState(), callArguments));
-    Set<DfaMemoryState> finalStates = ContainerUtil.newLinkedHashSet();
-    for (MethodContract contract : contracts) {
-      currentStates = addContractResults(contract, currentStates, factory, finalStates, defaultResult, methodRef);
-    }
-    return StreamEx.of(finalStates).map(DfaMemoryState::peek)
-      .append(currentStates.isEmpty() ? StreamEx.empty() : StreamEx.of(defaultResult)).distinct();
-  }
-
-  protected void processMethodReferenceResult(PsiMethodReferenceExpression methodRef,
-                                              List<? extends MethodContract> contracts,
-                                              DfaValue res) {
-  }
-
   @Override
   public DfaInstructionState[] visitTypeCast(TypeCastInstruction instruction, DataFlowRunner runner, DfaMemoryState memState) {
     PsiType type = instruction.getCastTo();
@@ -335,17 +318,12 @@ public class StandardInstructionVisitor extends InstructionVisitor {
       }
     }
 
-    PsiMethodReferenceExpression methodRef = instruction.getMethodType() == MethodCallInstruction.MethodType.METHOD_REFERENCE_CALL ?
-                                             (PsiMethodReferenceExpression)instruction.getContext() : null;
     DfaInstructionState[] result = new DfaInstructionState[finalStates.size()];
     int i = 0;
     for (DfaMemoryState state : finalStates) {
       if (instruction.shouldFlushFields()) {
         state.flushFields();
       }
-      if (methodRef != null) {
-        processMethodReferenceResult(methodRef, instruction.getContracts(), state.peek());
-      }
       result[i++] = new DfaInstructionState(runner.getInstruction(instruction.getIndex() + 1), state);
     }
     return result;