/// \brief If an overloaded operator call is a dereference of IndexVar or
/// a subscript of a the container with IndexVar as the single argument,
/// include it as a valid usage and prune the traversal.
///
/// For example, given
/// \code
///   struct Foo { int bar(); int x; };
///   vector<Foo> v;
///   void f(Foo);
/// \endcode
/// the following uses will be considered convertible:
/// \code
///   for (vector<Foo>::iterator i = v.begin(), e = v.end(); i != e; ++i) {
///     f(*i);
///   }
///   for (int i = 0; i < v.size(); ++i) {
///      int i = v[i] + 1;
///   }
/// \endcode
bool ForLoopIndexUseVisitor::TraverseCXXOperatorCallExpr(
    CXXOperatorCallExpr *OpCall) {
  switch (OpCall->getOperator()) {
  case OO_Star:
    if (isDereferenceOfOpCall(OpCall, IndexVar)) {
      Usages.push_back(Usage(OpCall));
      return true;
    }
    break;

  case OO_Subscript:
    if (OpCall->getNumArgs() != 2)
      break;
    if (isIndexInSubscriptExpr(Context, OpCall->getArg(1), IndexVar,
                               OpCall->getArg(0), ContainerExpr,
                               ContainerNeedsDereference)) {
      Usages.push_back(Usage(OpCall));
      return true;
    }
    break;

  default:
    break;
  }
  return VisitorBase::TraverseCXXOperatorCallExpr(OpCall);
}
//// \brief Apply the source transformations necessary to migrate the loop!
void LoopFixer::doConversion(ASTContext *Context,
                             const VarDecl *IndexVar,
                             const VarDecl *MaybeContainer,
                             StringRef ContainerString,
                             const UsageResult &Usages,
                             const DeclStmt *AliasDecl, const ForStmt *TheLoop,
                             bool ContainerNeedsDereference) {
  std::string VarName;

  if (Usages.size() == 1 && AliasDecl) {
    const VarDecl *AliasVar = cast<VarDecl>(AliasDecl->getSingleDecl());
    VarName = AliasVar->getName().str();
    // We keep along the entire DeclStmt to keep the correct range here.
    const SourceRange &ReplaceRange = AliasDecl->getSourceRange();
    if (!CountOnly)
      Replace->insert(
          Replacement(Context->getSourceManager(),
                      CharSourceRange::getTokenRange(ReplaceRange), ""));
    // No further replacements are made to the loop, since the iterator or index
    // was used exactly once - in the initialization of AliasVar.
  } else {
    VariableNamer Namer(GeneratedDecls, &ParentFinder->getStmtToParentStmtMap(),
                        TheLoop, IndexVar, MaybeContainer);
    VarName = Namer.createIndexName();
    // First, replace all usages of the array subscript expression with our new
    // variable.
    for (UsageResult::const_iterator I = Usages.begin(), E = Usages.end();
         I != E; ++I) {
      std::string ReplaceText = I->IsArrow ? VarName + "." : VarName;
      ReplacedVarRanges->insert(std::make_pair(TheLoop, IndexVar));
      if (!CountOnly)
        Replace->insert(
            Replacement(Context->getSourceManager(),
                        CharSourceRange::getTokenRange(I->Range),
                        ReplaceText));
    }
  }

  // Now, we need to construct the new range expresion.
  SourceRange ParenRange(TheLoop->getLParenLoc(), TheLoop->getRParenLoc());

  QualType AutoRefType =
      Context->getLValueReferenceType(Context->getAutoDeductType());

  std::string MaybeDereference = ContainerNeedsDereference ? "*" : "";
  std::string TypeString = AutoRefType.getAsString();
  std::string Range = ("(" + TypeString + " " + VarName + " : "
                           + MaybeDereference + ContainerString + ")").str();
  if (!CountOnly)
    Replace->insert(Replacement(Context->getSourceManager(),
                                CharSourceRange::getTokenRange(ParenRange),
                                Range));
  GeneratedDecls->insert(make_pair(TheLoop, VarName));
}
/// \brief If the unary operator is a dereference of IndexVar, include it
/// as a valid usage and prune the traversal.
///
/// For example, if container.begin() and container.end() both return pointers
/// to int, this makes sure that the initialization for `k` is not counted as an
/// unconvertible use of the iterator `i`.
/// \code
///   for (int *i = container.begin(), *e = container.end(); i != e; ++i) {
///     int k = *i + 2;
///   }
/// \endcode
bool ForLoopIndexUseVisitor::TraverseUnaryDeref(UnaryOperator *Uop) {
  // If we dereference an iterator that's actually a pointer, count the
  // occurrence.
  if (isDereferenceOfUop(Uop, IndexVar)) {
    Usages.push_back(Usage(Uop));
    return true;
  }

  return VisitorBase::TraverseUnaryOperator(Uop);
}
/// \brief If the member expression is operator-> (overloaded or not) on
/// IndexVar, include it as a valid usage and prune the traversal.
///
/// For example, given
/// \code
///   struct Foo { int bar(); int x; };
///   vector<Foo> v;
/// \endcode
/// the following uses will be considered convertible:
/// \code
///   for (vector<Foo>::iterator i = v.begin(), e = v.end(); i != e; ++i) {
///     int b = i->bar();
///     int k = i->x + 1;
///   }
/// \endcode
/// though
/// \code
///   for (vector<Foo>::iterator i = v.begin(), e = v.end(); i != e; ++i) {
///     int k = i.insert(1);
///   }
///   for (vector<Foo>::iterator i = v.begin(), e = v.end(); i != e; ++i) {
///     int b = e->bar();
///   }
/// \endcode
/// will not.
bool ForLoopIndexUseVisitor::TraverseMemberExpr(MemberExpr *Member) {
  const Expr *Base = Member->getBase();
  const DeclRefExpr *Obj = getDeclRef(Base);
  const Expr *ResultExpr = Member;
  QualType ExprType;
  if (const CXXOperatorCallExpr *Call =
      dyn_cast<CXXOperatorCallExpr>(Base->IgnoreParenImpCasts())) {
    // If operator->() is a MemberExpr containing a CXXOperatorCallExpr, then
    // the MemberExpr does not have the expression we want. We therefore catch
    // that instance here.
    // For example, if vector<Foo>::iterator defines operator->(), then the
    // example `i->bar()` at the top of this function is a CXXMemberCallExpr
    // referring to `i->` as the member function called. We want just `i`, so
    // we take the argument to operator->() as the base object.
    if(Call->getOperator() == OO_Arrow) {
      assert(Call->getNumArgs() == 1 &&
             "Operator-> takes more than one argument");
      Obj = getDeclRef(Call->getArg(0));
      ResultExpr = Obj;
      ExprType = Call->getCallReturnType();
    }
  }

  if (Member->isArrow() && Obj && exprReferencesVariable(IndexVar, Obj)) {
    if (ExprType.isNull())
      ExprType = Obj->getType();

    assert(ExprType->isPointerType() && "Operator-> returned non-pointer type");
    // FIXME: This works around not having the location of the arrow operator.
    // Consider adding OperatorLoc to MemberExpr?
    SourceLocation ArrowLoc =
        Lexer::getLocForEndOfToken(Base->getExprLoc(), 0,
                                   Context->getSourceManager(),
                                   Context->getLangOpts());
    // If something complicated is happening (i.e. the next token isn't an
    // arrow), give up on making this work.
    if (!ArrowLoc.isInvalid()) {
      Usages.push_back(Usage(ResultExpr, /*IsArrow=*/true,
                             SourceRange(Base->getExprLoc(), ArrowLoc)));
      return true;
    }
  }
  return TraverseStmt(Member->getBase());
}
/// \brief If we encounter an array with IndexVar as the index of an
/// ArraySubsriptExpression, note it as a consistent usage and prune the
/// AST traversal.
///
/// For example, given
/// \code
///   const int N = 5;
///   int arr[N];
/// \endcode
/// This is intended to permit
/// \code
///   for (int i = 0; i < N; ++i) {  /* use arr[i] */ }
/// \endcode
/// but not
/// \code
///   for (int i = 0; i < N; ++i) {  /* use notArr[i] */ }
/// \endcode
/// and further checking needs to be done later to ensure that exactly one array
/// is referenced.
bool ForLoopIndexUseVisitor::TraverseArraySubscriptExpr(
    ArraySubscriptExpr *ASE) {
  Expr *Arr = ASE->getBase();
  if (!isIndexInSubscriptExpr(ASE->getIdx(), IndexVar))
    return VisitorBase::TraverseArraySubscriptExpr(ASE);

  if ((ContainerExpr && !areSameExpr(Context, Arr->IgnoreParenImpCasts(),
                                     ContainerExpr->IgnoreParenImpCasts()))
      || !arrayMatchesBoundExpr(Context, Arr->IgnoreImpCasts()->getType(),
                                ArrayBoundExpr)) {
    // If we have already discovered the array being indexed and this isn't it
    // or this array doesn't match, mark this loop as unconvertible.
    OnlyUsedAsIndex = false;
    return VisitorBase::TraverseArraySubscriptExpr(ASE);
  }

  if (!ContainerExpr)
    ContainerExpr = Arr;

  Usages.push_back(Usage(ASE));
  return true;
}
/// \brief If a member function call is the at() accessor on the container with
/// IndexVar as the single argument, include it as a valid usage and prune
/// the traversal.
///
/// Member calls on other objects will not be permitted.
/// Calls on the iterator object are not permitted, unless done through
/// operator->(). The one exception is allowing vector::at() for pseudoarrays.
bool ForLoopIndexUseVisitor::TraverseCXXMemberCallExpr(
    CXXMemberCallExpr *MemberCall) {
  MemberExpr *Member =
      dyn_cast<MemberExpr>(MemberCall->getCallee()->IgnoreParenImpCasts());
  if (!Member)
    return VisitorBase::TraverseCXXMemberCallExpr(MemberCall);
  // We specifically allow an accessor named "at" to let STL in, though
  // this is restricted to pseudo-arrays by requiring a single, integer
  // argument.
  const IdentifierInfo *Ident = Member->getMemberDecl()->getIdentifier();
  if (Ident && Ident->isStr("at") && MemberCall->getNumArgs() == 1) {
    if (isIndexInSubscriptExpr(Context, MemberCall->getArg(0), IndexVar,
                               Member->getBase(), ContainerExpr,
                               ContainerNeedsDereference)) {
      Usages.push_back(Usage(MemberCall));
      return true;
    }
  }

  if (containsExpr(Context, &DependentExprs, Member->getBase()))
    ConfidenceLevel.lowerTo(TCK_Risky);

  return VisitorBase::TraverseCXXMemberCallExpr(MemberCall);
}
//// \brief Apply the source transformations necessary to migrate the loop!
void LoopFixer::doConversion(ASTContext *Context,
                             const VarDecl *IndexVar,
                             const VarDecl *MaybeContainer,
                             StringRef ContainerString,
                             const UsageResult &Usages,
                             const DeclStmt *AliasDecl,
                             bool AliasUseRequired,
                             bool AliasFromForInit,
                             const ForStmt *TheLoop,
                             bool ContainerNeedsDereference,
                             bool DerefByValue,
                             bool DerefByConstRef) {
  std::string VarName;
  bool VarNameFromAlias = Usages.size() == 1 && AliasDecl;
  bool AliasVarIsRef = false;

  if (VarNameFromAlias) {
    const VarDecl *AliasVar = cast<VarDecl>(AliasDecl->getSingleDecl());
    VarName = AliasVar->getName().str();
    AliasVarIsRef = AliasVar->getType()->isReferenceType();

    // We keep along the entire DeclStmt to keep the correct range here.
    const SourceRange &ReplaceRange = AliasDecl->getSourceRange();

    std::string ReplacementText;
    if (AliasUseRequired)
      ReplacementText = VarName;
    else if (AliasFromForInit)
      // FIXME: Clang includes the location of the ';' but only for DeclStmt's
      // in a for loop's init clause. Need to put this ';' back while removing
      // the declaration of the alias variable. This is probably a bug.
      ReplacementText = ";";

    Owner.addReplacementForCurrentTU(Replacement(
        Context->getSourceManager(),
        CharSourceRange::getTokenRange(ReplaceRange), ReplacementText));
    // No further replacements are made to the loop, since the iterator or index
    // was used exactly once - in the initialization of AliasVar.
  } else {
    VariableNamer Namer(&TUInfo.getGeneratedDecls(),
                        &TUInfo.getParentFinder().getStmtToParentStmtMap(),
                        TheLoop, IndexVar, MaybeContainer, Context);
    VarName = Namer.createIndexName();
    // First, replace all usages of the array subscript expression with our new
    // variable.
    for (UsageResult::const_iterator I = Usages.begin(), E = Usages.end();
         I != E; ++I) {
      std::string ReplaceText = I->IsArrow ? VarName + "." : VarName;
      TUInfo.getReplacedVars().insert(std::make_pair(TheLoop, IndexVar));
      Owner.addReplacementForCurrentTU(
          Replacement(Context->getSourceManager(),
                      CharSourceRange::getTokenRange(I->Range), ReplaceText));
    }
  }

  // Now, we need to construct the new range expresion.
  SourceRange ParenRange(TheLoop->getLParenLoc(), TheLoop->getRParenLoc());

  QualType AutoRefType = Context->getAutoDeductType();

  // If the new variable name is from the aliased variable, then the reference
  // type for the new variable should only be used if the aliased variable was
  // declared as a reference.
  if (!VarNameFromAlias || AliasVarIsRef) {
    // If an iterator's operator*() returns a 'T&' we can bind that to 'auto&'.
    // If operator*() returns 'T' we can bind that to 'auto&&' which will deduce
    // to 'T&&'.
    if (DerefByValue)
      AutoRefType = Context->getRValueReferenceType(AutoRefType);
    else {
      if (DerefByConstRef)
        AutoRefType = Context->getConstType(AutoRefType);
      AutoRefType = Context->getLValueReferenceType(AutoRefType);
    }
  }

  std::string MaybeDereference = ContainerNeedsDereference ? "*" : "";
  std::string TypeString = AutoRefType.getAsString();
  std::string Range = ("(" + TypeString + " " + VarName + " : "
                           + MaybeDereference + ContainerString + ")").str();
  Owner.addReplacementForCurrentTU(
      Replacement(Context->getSourceManager(),
                  CharSourceRange::getTokenRange(ParenRange), Range));
  TUInfo.getGeneratedDecls().insert(make_pair(TheLoop, VarName));
}