GuessManagerImpl: support type constrains from DFA, various improvements

Partially fixes IDEA-181794 Suggest completion items with type casts in stream chains if there exists 'filter(x -> x instanceof Foo)' call
Fixes IDEA-182455 Code completion: resolve actual value type assigned to local variable

java/java-analysis-impl/src/com/intellij/codeInsight/guess/impl/ExpressionTypeMemoryState.java

@@ -75,6 +75,7 @@ public class ExpressionTypeMemoryState extends DfaMemoryStateImpl {
       if (!value.isNegated()) {
         setExpressionType(value.getExpression(), value.getCastType());
       }
+      return super.applyCondition(((DfaInstanceofValue)dfaCond).getRelation());
     }
 
     return super.applyCondition(dfaCond);

java/java-analysis-impl/src/com/intellij/codeInsight/guess/impl/GuessManagerImpl.java

@@ -19,6 +19,8 @@ import com.intellij.codeInsight.guess.GuessManager;
 import com.intellij.codeInspection.dataFlow.*;
 import com.intellij.codeInspection.dataFlow.instructions.*;
 import com.intellij.codeInspection.dataFlow.value.DfaInstanceofValue;
+import com.intellij.codeInspection.dataFlow.value.DfaRelationValue;
+import com.intellij.codeInspection.dataFlow.value.DfaValue;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
@@ -33,7 +35,7 @@ import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.BitUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
-import gnu.trove.THashMap;
+import com.siyeh.ig.psiutils.ExpressionUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -153,33 +155,60 @@ public class GuessManagerImpl extends GuessManager {
     };
 
     final ExpressionTypeInstructionVisitor visitor = new ExpressionTypeInstructionVisitor(forPlace);
-    if (runner.analyzeMethod(scope, visitor) == RunnerResult.OK) {
+    if (runner.analyzeMethodWithInlining(scope, visitor) == RunnerResult.OK) {
       return visitor.getResult();
     }
     return null;
   }
 
-  private static Map<PsiExpression, PsiType> getAllTypeCasts(PsiExpression forPlace) {
+  private static boolean mayHaveMorePreciseType(PsiExpression expr) {
-    assert forPlace.isValid();
+    PsiExpression place = PsiUtil.skipParenthesizedExprDown(expr);
-    final int start = forPlace.getTextRange().getStartOffset();
+    if (place instanceof PsiReferenceExpression) {
-    final Map<PsiExpression, PsiType> allCasts = new THashMap<>(ExpressionTypeMemoryState.EXPRESSION_HASHING_STRATEGY);
+      PsiElement target = ((PsiReferenceExpression)place).resolve();
-    getTopmostBlock(forPlace).accept(new JavaRecursiveElementWalkingVisitor() {
+      if (target instanceof PsiParameter) {
+        PsiElement parent = target.getParent();
+        if (parent instanceof PsiParameterList && parent.getParent() instanceof PsiLambdaExpression) {
+          return true;
+        }
+      }
+    }
+    if (place == null) return false;
+    final int start = place.getTextRange().getStartOffset();
+    class Visitor extends JavaRecursiveElementWalkingVisitor {
+      public boolean hasInteresting;
+
+      @Override
+      public void visitAssignmentExpression(PsiAssignmentExpression expression) {
+        if (ExpressionTypeMemoryState.EXPRESSION_HASHING_STRATEGY.equals(expression.getLExpression(), place)) {
+          hasInteresting = true;
+          stopWalking();
+        }
+        super.visitAssignmentExpression(expression);
+      }
+
+      @Override
+      public void visitLocalVariable(PsiLocalVariable variable) {
+        if (variable.getInitializer() != null && ExpressionUtils.isReferenceTo(place, variable)) {
+          hasInteresting = true;
+          stopWalking();
+        }
+        super.visitLocalVariable(variable);
+      }
+
       @Override
       public void visitTypeCastExpression(PsiTypeCastExpression expression) {
-        final PsiType castType = expression.getType();
+        if (ExpressionTypeMemoryState.EXPRESSION_HASHING_STRATEGY.equals(expression.getOperand(), place)) {
-        final PsiExpression operand = expression.getOperand();
+          hasInteresting = true;
-        if (operand != null && castType != null) {
+          stopWalking();
-          allCasts.put(operand, castType);
         }
         super.visitTypeCastExpression(expression);
       }
 
       @Override
       public void visitInstanceOfExpression(PsiInstanceOfExpression expression) {
-        final PsiTypeElement castType = expression.getCheckType();
+        if (ExpressionTypeMemoryState.EXPRESSION_HASHING_STRATEGY.equals(expression.getOperand(), place)) {
-        final PsiExpression operand = expression.getOperand();
+          hasInteresting = true;
-        if (castType != null) {
+          stopWalking();
-          allCasts.put(operand, castType.getType());
         }
         super.visitInstanceOfExpression(expression);
       }
@@ -187,13 +216,14 @@ public class GuessManagerImpl extends GuessManager {
       @Override
       public void visitElement(PsiElement element) {
         if (element.getTextRange().getStartOffset() > start) {
-          return;
+          stopWalking();
         }
-
         super.visitElement(element);
       }
-    });
+    }
-    return allCasts;
+    Visitor visitor = new Visitor();
+    getTopmostBlock(place).accept(visitor);
+    return visitor.hasInteresting;
   }
 
   private static PsiElement getTopmostBlock(PsiElement scope) {
@@ -372,8 +402,7 @@ public class GuessManagerImpl extends GuessManager {
   @NotNull
   @Override
   public List<PsiType> getControlFlowExpressionTypeConjuncts(@NotNull PsiExpression expr) {
-    final Map<PsiExpression, PsiType> allCasts = getAllTypeCasts(expr);
+    if (!mayHaveMorePreciseType(expr)) {
-    if (!allCasts.containsKey(expr)) {
       return Collections.emptyList(); //optimization
     }
 
@@ -388,9 +417,10 @@ public class GuessManagerImpl extends GuessManager {
     return Collections.emptyList();
   }
 
-  private static class ExpressionTypeInstructionVisitor extends InstructionVisitor {
+  private static class ExpressionTypeInstructionVisitor extends StandardInstructionVisitor {
     private MultiMap<PsiExpression, PsiType> myResult;
     private final PsiElement myForPlace;
+    private TypeConstraint myConstraint = null;
 
     private ExpressionTypeInstructionVisitor(@NotNull PsiElement forPlace) {
       PsiElement parent = PsiUtil.skipParenthesizedExprUp(forPlace.getParent());
@@ -402,14 +432,24 @@ public class GuessManagerImpl extends GuessManager {
     }
 
     MultiMap<PsiExpression, PsiType> getResult() {
+      if (myConstraint != null && myForPlace instanceof PsiExpression) {
+        PsiType type = myConstraint.getPsiType();
+        if (type instanceof PsiIntersectionType) {
+          myResult.putValues((PsiExpression)myForPlace, Arrays.asList(((PsiIntersectionType)type).getConjuncts()));
+        }
+        else {
+          myResult.putValue((PsiExpression)myForPlace, type);
+        }
+      }
       return myResult;
     }
 
     @Override
     public DfaInstructionState[] visitInstanceof(InstanceofInstruction instruction, DataFlowRunner runner, DfaMemoryState memState) {
-      memState.pop();
+      DfaValue type = memState.pop();
-      memState.pop();
+      DfaValue operand = memState.pop();
-      memState.push(new DfaInstanceofValue(runner.getFactory(), instruction.getLeft(), instruction.getCastType()));
+      DfaValue relation = runner.getFactory().createCondition(operand, DfaRelationValue.RelationType.IS, type);
+      memState.push(new DfaInstanceofValue(runner.getFactory(), instruction.getLeft(), instruction.getCastType(), relation, false));
       return new DfaInstructionState[]{new DfaInstructionState(runner.getInstruction(instruction.getIndex() + 1), memState)};
     }
 
@@ -426,10 +466,6 @@ public class GuessManagerImpl extends GuessManager {
       if (left != null && right != null) {
         MultiMap<PsiExpression, PsiType> states = ((ExpressionTypeMemoryState)memState).getStates();
         states.remove(left);
-        PsiType rightType = right.getType();
-        if (rightType != null) {
-          ((ExpressionTypeMemoryState) memState).setExpressionType(left, runner.getFactory().createDfaType(rightType).getPsiType());
-        }
       }
       return super.visitAssign(instruction, runner, memState);
     }
@@ -439,7 +475,11 @@ public class GuessManagerImpl extends GuessManager {
       if (myForPlace == instruction.getCallExpression()) {
         addToResult(((ExpressionTypeMemoryState)memState).getStates());
       }
-      return super.visitMethodCall(instruction, runner, memState);
+      DfaInstructionState[] states = super.visitMethodCall(instruction, runner, memState);
+      if (myForPlace == instruction.getCallExpression()) {
+        addConstraints(states);
+      }
+      return states;
     }
 
     @Override
@@ -447,7 +487,22 @@ public class GuessManagerImpl extends GuessManager {
       if (myForPlace == instruction.getPlace()) {
         addToResult(((ExpressionTypeMemoryState)memState).getStates());
       }
-      return super.visitPush(instruction, runner, memState);
+      DfaInstructionState[] states = super.visitPush(instruction, runner, memState);
+      if (myForPlace == instruction.getPlace()) {
+        addConstraints(states);
+      }
+      return states;
+    }
+
+    private void addConstraints(DfaInstructionState[] states) {
+      for (DfaInstructionState state : states) {
+        DfaMemoryState memoryState = state.getMemoryState();
+        if (myConstraint == TypeConstraint.EMPTY) return;
+        TypeConstraint constraint = memoryState.getValueFact(memoryState.peek(), DfaFactType.TYPE_CONSTRAINT);
+        if (constraint != null) {
+          myConstraint = myConstraint == null ? constraint : myConstraint.union(constraint);
+        }
+      }
     }
 
     private void addToResult(MultiMap<PsiExpression, PsiType> map) {

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

@@ -76,14 +76,16 @@ public class DataFlowRunner {
   }
 
   @Nullable
-  private Collection<DfaMemoryState> createInitialStates(@NotNull PsiElement psiBlock, @NotNull InstructionVisitor visitor) {
+  private Collection<DfaMemoryState> createInitialStates(@NotNull PsiElement psiBlock,
+                                                         @NotNull InstructionVisitor visitor,
+                                                         boolean allowInlining) {
     PsiElement container = PsiTreeUtil.getParentOfType(psiBlock, PsiClass.class, PsiLambdaExpression.class);
     if (container != null && (!(container instanceof PsiClass) || PsiUtil.isLocalOrAnonymousClass((PsiClass)container))) {
       PsiElement block = DfaPsiUtil.getTopmostBlockInSameClass(container.getParent());
       if (block != null) {
         final RunnerResult result;
         try {
-          myInlining = false;
+          myInlining = allowInlining;
           result = analyzeMethod(block, visitor);
         }
         finally {
@@ -91,7 +93,7 @@ public class DataFlowRunner {
         }
         if (result == RunnerResult.OK) {
           final Collection<DfaMemoryState> closureStates = myNestedClosures.get(DfaPsiUtil.getTopmostBlockInSameClass(psiBlock));
-          if (!closureStates.isEmpty()) {
+          if (allowInlining || !closureStates.isEmpty()) {
             return closureStates;
           }
         }
@@ -102,12 +104,41 @@ public class DataFlowRunner {
     return Collections.singletonList(createMemoryState());
   }
 
+  /**
+   * Analyze this particular method (lambda, class initializer) without inlining this method into parent one.
+   * E.g. if supplied method is a lambda within Stream API call chain, it still will be analyzed as separate method.
+   * On the other hand, inlining will normally work inside the supplied method.
+   *
+   * @param psiBlock method/lambda/class initializer body
+   * @param visitor a visitor to use
+   * @return result status
+   */
   @NotNull
   public final RunnerResult analyzeMethod(@NotNull PsiElement psiBlock, @NotNull InstructionVisitor visitor) {
-    Collection<DfaMemoryState> initialStates = createInitialStates(psiBlock, visitor);
+    Collection<DfaMemoryState> initialStates = createInitialStates(psiBlock, visitor, false);
     return initialStates == null ? RunnerResult.NOT_APPLICABLE : analyzeMethod(psiBlock, visitor, false, initialStates);
   }
 
+  /**
+   * Analyze this particular method (lambda, class initializer) trying to inline it into outer scope if possible.
+   * Usually inlining works, e.g. for lambdas inside stream API calls.
+   *
+   * @param psiBlock method/lambda/class initializer body
+   * @param visitor a visitor to use
+   * @return result status
+   */
+  @NotNull
+  public final RunnerResult analyzeMethodWithInlining(@NotNull PsiElement psiBlock, @NotNull InstructionVisitor visitor) {
+    Collection<DfaMemoryState> initialStates = createInitialStates(psiBlock, visitor, true);
+    if (initialStates == null) {
+      return RunnerResult.NOT_APPLICABLE;
+    }
+    if (initialStates.isEmpty()) {
+      return RunnerResult.OK;
+    }
+    return analyzeMethod(psiBlock, visitor, false, initialStates);
+  }
+
   public final RunnerResult analyzeCodeBlock(@NotNull PsiCodeBlock block,
                                              @NotNull InstructionVisitor visitor,
                                              Consumer<DfaMemoryState> initialStateAdjuster) {
@@ -247,7 +278,7 @@ public class DataFlowRunner {
   }
 
   public RunnerResult analyzeMethodRecursively(PsiElement block, StandardInstructionVisitor visitor) {
-    Collection<DfaMemoryState> states = createInitialStates(block, visitor);
+    Collection<DfaMemoryState> states = createInitialStates(block, visitor, false);
     if (states == null) return RunnerResult.NOT_APPLICABLE;
     return analyzeBlockRecursively(block, states, visitor);
   }

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

@@ -24,10 +24,7 @@ import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.UnorderedPair;
 import com.intellij.openapi.util.text.StringUtil;
-import com.intellij.psi.PsiModifierListOwner;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiPrimitiveType;
-import com.intellij.psi.PsiType;
-import com.intellij.psi.PsiVariable;
 import com.intellij.psi.util.TypeConversionUtil;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.ObjectUtils;
@@ -904,12 +901,19 @@ public class DfaMemoryStateImpl implements DfaMemoryState {
           setVariableState(dfaVar, getVariableState(dfaVar).withFact(DfaFactType.CAN_BE_NULL, true));
           return;
         }
+        PsiType psiType;
         if (constValue instanceof PsiVariable) {
-          DfaPsiType dfaType = myFactory.createDfaType(((PsiVariable)constValue).getType());
+          psiType = ((PsiVariable)constValue).getType();
-          DfaVariableState state = getVariableState(dfaVar).withInstanceofValue(dfaType);
+        }
-          if (state != null) {
+        else {
-            setVariableState(dfaVar, state);
+          PsiModifierListOwner context = dfaVar.getPsiVariable();
-          }
+          psiType = JavaPsiFacade.getElementFactory(context.getProject())
+            .createTypeByFQClassName(constValue.getClass().getName(), context.getResolveScope());
+        }
+        DfaPsiType dfaType = myFactory.createDfaType(psiType);
+        DfaVariableState state = getVariableState(dfaVar).withInstanceofValue(dfaType);
+        if (state != null) {
+          setVariableState(dfaVar, state);
         }
       }
       if (isNotNull(value) && !isNotNull(dfaVar)) {

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

@@ -162,7 +162,7 @@ public final class TypeConstraint {
   }
 
   @Nullable
-  TypeConstraint union(@NotNull TypeConstraint other) {
+  public TypeConstraint union(@NotNull TypeConstraint other) {
     if(isSuperStateOf(other)) return this;
     if(other.isSuperStateOf(this)) return other;
     Set<DfaPsiType> leftTypes = new HashSet<>(this.myInstanceofValues);

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

@@ -26,18 +26,25 @@ public class DfaInstanceofValue extends DfaValue {
   private final @NotNull PsiExpression myExpression;
   private final @NotNull PsiType myCastType;
   private final boolean myNegated;
+  private final @NotNull DfaValue myRelation;
 
-  public DfaInstanceofValue(DfaValueFactory factory, @NotNull PsiExpression expression, @NotNull PsiType castType) {
+  public DfaInstanceofValue(DfaValueFactory factory,
-    this(factory, expression, castType, false);
+                            @NotNull PsiExpression expression,
-  }
+                            @NotNull PsiType castType,
-
+                            @NotNull DfaValue relation,
-  public DfaInstanceofValue(DfaValueFactory factory, @NotNull PsiExpression expression, @NotNull PsiType castType, boolean negated) {
+                            boolean negated) {
     super(factory);
     myExpression = expression;
     myCastType = castType;
+    myRelation = relation;
     myNegated = negated;
   }
 
+  @NotNull
+  public DfaValue getRelation() {
+    return myRelation;
+  }
+
   @NotNull
   public PsiExpression getExpression() {
     return myExpression;
@@ -54,6 +61,6 @@ public class DfaInstanceofValue extends DfaValue {
 
   @Override
   public DfaValue createNegated() {
-    return new DfaInstanceofValue(myFactory, myExpression, myCastType, !myNegated);
+    return new DfaInstanceofValue(myFactory, myExpression, myCastType, myRelation.createNegated(), !myNegated);
   }
 }

java/java-tests/testData/codeInsight/completion/normal/AssignmentPreciseTypeDfa.java

@@ -0,0 +1,7 @@
+class Foo {
+    void test(String s) {
+        Object x;
+        x = s;
+        System.out.println(x.subst<caret>);
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/AssignmentPreciseTypeDfa_after.java

@@ -0,0 +1,7 @@
+class Foo {
+    void test(String s) {
+        Object x;
+        x = s;
+        System.out.println(((String) x).substring());
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/DeclarationPreciseTypeDfa.java

@@ -0,0 +1,6 @@
+class Foo {
+    void test(String s) {
+        Object x = s;
+        System.out.println(x.subst<caret>);
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/DeclarationPreciseTypeDfa_after.java

@@ -0,0 +1,6 @@
+class Foo {
+    void test(String s) {
+        Object x = s;
+        System.out.println(((String) x).substring());
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/OptionalDfa.java

@@ -0,0 +1,8 @@
+import java.util.*;
+
+class Foo {
+    void test(Optional<Object> opt) {
+        opt.filter(x -> x instanceof String)
+          .map(s -> s.subst<caret>)
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/OptionalDfa_after.java

@@ -0,0 +1,8 @@
+import java.util.*;
+
+class Foo {
+    void test(Optional<Object> opt) {
+        opt.filter(x -> x instanceof String)
+          .map(s -> ((String) s).substring())
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/StreamDfa.java

@@ -0,0 +1,9 @@
+import java.util.*;
+
+class Foo {
+    void test(List<?> obj) {
+        obj.stream()
+          .filter(x -> x instanceof String)
+          .forEach(e -> e.subst<caret>);
+    }
+}

java/java-tests/testData/codeInsight/completion/normal/StreamDfa_after.java

@@ -0,0 +1,9 @@
+import java.util.*;
+
+class Foo {
+    void test(List<?> obj) {
+        obj.stream()
+          .filter(x -> x instanceof String)
+          .forEach(e -> ((String) e).substring());
+    }
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/NormalCompletionDfaTest.groovy

@@ -17,6 +17,7 @@ package com.intellij.java.codeInsight.completion
 
 import com.intellij.JavaTestUtil
 import com.intellij.codeInsight.completion.LightFixtureCompletionTestCase
+import com.intellij.testFramework.LightProjectDescriptor
 
 /**
  * @author peter
@@ -26,7 +27,12 @@ class NormalCompletionDfaTest extends LightFixtureCompletionTestCase {
   protected String getBasePath() {
     return JavaTestUtil.getRelativeJavaTestDataPath() + "/codeInsight/completion/normal/"
   }
-  
+
+  @Override
+  protected LightProjectDescriptor getProjectDescriptor() {
+    return JAVA_8
+  }
+
   void testCastInstanceofedQualifier() { doTest() }
   void testCastInstanceofedQualifierInForeach() { doTest() }
   void testCastComplexInstanceofedQualifier() { doTest() }
@@ -42,7 +48,11 @@ class NormalCompletionDfaTest extends LightFixtureCompletionTestCase {
   void testQualifierCastingBeforeLt() { doTest() }
   void testCastQualifierForPrivateFieldReference() { doTest() }
   void testOrAssignmentDfa() { doTest() }
+  void testAssignmentPreciseTypeDfa() { doTest() }
+  void testDeclarationPreciseTypeDfa() { doTest() }
   void testInstanceOfAssignmentDfa() { doTest() }
+  void testStreamDfa() { doTest() }
+  void testOptionalDfa() { doTest() }
   void testFieldWithCastingCaret() { doTest() }
 
   void testCastTwice() {