@Override
  public PsiType[] guessTypeToCast(
      PsiExpression expr) { // TODO : make better guess based on control flow
    LinkedHashSet<PsiType> types = new LinkedHashSet<>();

    ContainerUtil.addIfNotNull(types, getControlFlowExpressionType(expr));
    addExprTypesWhenContainerElement(types, expr);
    addExprTypesByDerivedClasses(types, expr);

    return types.toArray(PsiType.createArray(types.size()));
  }
  private void addExprTypesByDerivedClasses(LinkedHashSet<PsiType> set, PsiExpression expr) {
    PsiType type = expr.getType();
    if (!(type instanceof PsiClassType)) return;
    PsiClass refClass = PsiUtil.resolveClassInType(type);
    if (refClass == null) return;

    PsiManager manager = PsiManager.getInstance(myProject);
    PsiElementProcessor.CollectElementsWithLimit<PsiClass> processor =
        new PsiElementProcessor.CollectElementsWithLimit<>(5);
    ClassInheritorsSearch.search(refClass).forEach(new PsiElementProcessorAdapter<>(processor));
    if (processor.isOverflow()) return;

    for (PsiClass derivedClass : processor.getCollection()) {
      if (derivedClass instanceof PsiAnonymousClass) continue;
      PsiType derivedType =
          JavaPsiFacade.getInstance(manager.getProject())
              .getElementFactory()
              .createType(derivedClass);
      set.add(derivedType);
    }
  }
 private void addExprTypesWhenContainerElement(LinkedHashSet<PsiType> set, PsiExpression expr) {
   if (expr instanceof PsiMethodCallExpression) {
     PsiMethodCallExpression callExpr = (PsiMethodCallExpression) expr;
     PsiReferenceExpression methodExpr = callExpr.getMethodExpression();
     String methodName = methodExpr.getReferenceName();
     MethodPattern pattern =
         myMethodPatternMap.findPattern(
             methodName, callExpr.getArgumentList().getExpressions().length);
     if (pattern != null && pattern.parameterIndex < 0 /* return value */) {
       PsiExpression qualifier = methodExpr.getQualifierExpression();
       if (qualifier != null) {
         PsiType[] types = guessContainerElementType(qualifier, null);
         for (PsiType type : types) {
           if (type instanceof PsiClassType) {
             if (((PsiClassType) type).resolve() instanceof PsiAnonymousClass) continue;
           }
           set.add(type);
         }
       }
     }
   }
 }