Stop LoopConvert removing DeclStmts from selection/iteration condition clauses

If the LoopConvert Transform detects an alias for the loop variable, it
attempts to use that name in the resulting range-based for loop while removing
the original DeclStmt for the variable. That removal produced bad code when the
declaration was in the condition of an if, switch, while, or for stmt. This
revision fixes the problem by simply replacing the declaration with a use of
the alias variable.

llvm-svn: 181242

3 files changed, 121 insertions, 10 deletions

diff --git a/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.cpp b/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.cpp
index 0647387..4ba22b3 100644
--- a/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.cpp
+++ b/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.cpp
@@ -72,7 +72,9 @@ class ForLoopIndexUseVisitor
     Context(Context), IndexVar(IndexVar), EndVar(EndVar),
     ContainerExpr(ContainerExpr), ArrayBoundExpr(ArrayBoundExpr),
     ContainerNeedsDereference(ContainerNeedsDereference),
-    OnlyUsedAsIndex(true),  AliasDecl(NULL), ConfidenceLevel(RL_Safe) {
+    OnlyUsedAsIndex(true),  AliasDecl(NULL), ConfidenceLevel(RL_Safe),
+    NextStmtParent(NULL), CurrStmtParent(NULL), ReplaceWithAliasUse(false),
+    AliasFromForInit(false) {
      if (ContainerExpr) {
        addComponent(ContainerExpr);
        llvm::FoldingSetNodeID ID;
@@ -123,6 +125,17 @@ class ForLoopIndexUseVisitor
     return ConfidenceLevel.getRiskLevel();
   }
 
+  /// \brief Indicates if the alias declaration was in a place where it cannot
+  /// simply be removed but rather replaced with a use of the alias variable.
+  /// For example, variables declared in the condition of an if, switch, or for
+  /// stmt.
+  bool aliasUseRequired() const { return ReplaceWithAliasUse; }
+
+  /// \brief Indicates if the alias declaration came from the init clause of a
+  /// nested for loop. SourceRanges provided by Clang for DeclStmts in this
+  /// case need to be adjusted.
+  bool aliasFromForInit() const { return AliasFromForInit; }
+
  private:
   /// Typedef used in CRTP functions.
   typedef RecursiveASTVisitor<ForLoopIndexUseVisitor> VisitorBase;
@@ -136,6 +149,7 @@ class ForLoopIndexUseVisitor
   bool TraverseUnaryDeref(UnaryOperator *Uop);
   bool VisitDeclRefExpr(DeclRefExpr *E);
   bool VisitDeclStmt(DeclStmt *S);
+  bool TraverseStmt(Stmt *S);
 
   /// \brief Add an expression to the list of expressions on which the container
   /// expression depends.
@@ -172,6 +186,19 @@ class ForLoopIndexUseVisitor
   /// of the loop element, lower our confidence level.
   llvm::SmallVector<
       std::pair<const Expr *, llvm::FoldingSetNodeID>, 16> DependentExprs;
+
+  /// The parent-in-waiting. Will become the real parent once we traverse down
+  /// one level in the AST.
+  const Stmt *NextStmtParent;
+  /// The actual parent of a node when Visit*() calls are made. Only the
+  /// parentage of DeclStmt's to possible iteration/selection statements is of
+  /// importance.
+  const Stmt *CurrStmtParent;
+
+  /// \see aliasUseRequired().
+  bool ReplaceWithAliasUse;
+  /// \see aliasFromForInit().
+  bool AliasFromForInit;
 };
 
 /// \brief Obtain the original source code text from a SourceRange.
@@ -722,18 +749,47 @@ bool ForLoopIndexUseVisitor::VisitDeclRefExpr(DeclRefExpr *E) {
 /// See the comments for isAliasDecl.
 bool ForLoopIndexUseVisitor::VisitDeclStmt(DeclStmt *S) {
   if (!AliasDecl && S->isSingleDecl() &&
-      isAliasDecl(S->getSingleDecl(), IndexVar))
-      AliasDecl = S;
+      isAliasDecl(S->getSingleDecl(), IndexVar)) {
+    AliasDecl = S;
+    if (CurrStmtParent) {
+      if (isa<IfStmt>(CurrStmtParent) ||
+          isa<WhileStmt>(CurrStmtParent) ||
+          isa<SwitchStmt>(CurrStmtParent))
+        ReplaceWithAliasUse = true;
+      else if (isa<ForStmt>(CurrStmtParent)) {
+        if (cast<ForStmt>(CurrStmtParent)->getConditionVariableDeclStmt() == S)
+          ReplaceWithAliasUse = true;
+        else
+          // It's assumed S came the for loop's init clause.
+          AliasFromForInit = true;
+      }
+    }
+  }
+
   return true;
 }
 
+bool ForLoopIndexUseVisitor::TraverseStmt(Stmt *S) {
+  // All this pointer swapping is a mechanism for tracking immediate parentage
+  // of Stmts.
+  const Stmt *OldNextParent = NextStmtParent;
+  CurrStmtParent = NextStmtParent;
+  NextStmtParent = S;
+  bool Result = VisitorBase::TraverseStmt(S);
+  NextStmtParent = OldNextParent;
+  return Result;
+}
+
 //// \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,
+                             const DeclStmt *AliasDecl,
+                             bool AliasUseRequired,
+                             bool AliasFromForInit,
+                             const ForStmt *TheLoop,
                              bool ContainerNeedsDereference,
                              bool DerefByValue) {
   std::string VarName;
@@ -747,9 +803,19 @@ void LoopFixer::doConversion(ASTContext *Context,
 
     // We keep along the entire DeclStmt to keep the correct range here.
     const SourceRange &ReplaceRange = AliasDecl->getSourceRange();
-    Replace->insert(
-        Replacement(Context->getSourceManager(),
-                    CharSourceRange::getTokenRange(ReplaceRange), ""));
+
+    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 = ";";
+
+    Replace->insert(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 {
@@ -945,9 +1011,9 @@ void LoopFixer::findAndVerifyUsages(ASTContext *Context,
     return;
 
   doConversion(Context, LoopVar, getReferencedVariable(ContainerExpr),
-               ContainerString, Finder.getUsages(),
-               Finder.getAliasDecl(), TheLoop, ContainerNeedsDereference,
-               DerefByValue);
+               ContainerString, Finder.getUsages(), Finder.getAliasDecl(),
+               Finder.aliasUseRequired(), Finder.aliasFromForInit(), TheLoop,
+               ContainerNeedsDereference, DerefByValue);
   ++*AcceptedChanges;
 }
 
diff --git a/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.h b/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.h
index 6f59c01..6a48889 100644
--- a/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.h
+++ b/clang-tools-extra/cpp11-migrate/LoopConvert/LoopActions.h
@@ -73,6 +73,8 @@ class LoopFixer : public clang::ast_matchers::MatchFinder::MatchCallback {
                     llvm::StringRef ContainerString,
                     const UsageResult &Usages,
                     const clang::DeclStmt *AliasDecl,
+                    bool AliasUseRequired,
+                    bool AliasFromForInit,
                     const clang::ForStmt *TheLoop,
                     bool ContainerNeedsDereference,
                     bool DerefByValue);
diff --git a/clang-tools-extra/test/cpp11-migrate/LoopConvert/naming-alias.cpp b/clang-tools-extra/test/cpp11-migrate/LoopConvert/naming-alias.cpp
index 0373d2b..0ed3440 100644
--- a/clang-tools-extra/test/cpp11-migrate/LoopConvert/naming-alias.cpp
+++ b/clang-tools-extra/test/cpp11-migrate/LoopConvert/naming-alias.cpp
@@ -8,6 +8,7 @@ const int N = 10;
 
 Val Arr[N];
 Val &func(Val &);
+void sideEffect(int);
 
 void aliasing() {
   // If the loop container is only used for a declaration of a temporary
@@ -56,6 +57,48 @@ void aliasing() {
   // CHECK: for (auto & elem : Arr)
   // CHECK-NEXT: Val &t = func(elem);
   // CHECK-NEXT: int y = t.x;
+
+  int IntArr[N];
+  for (unsigned i = 0; i < N; ++i) {
+    if (int alias = IntArr[i]) {
+      sideEffect(alias);
+    }
+  }
+  // CHECK: for (auto alias : IntArr)
+  // CHECK-NEXT: if (alias) {
+
+  for (unsigned i = 0; i < N; ++i) {
+    while (int alias = IntArr[i]) {
+      sideEffect(alias);
+    }
+  }
+  // CHECK: for (auto alias : IntArr)
+  // CHECK-NEXT: while (alias) {
+
+  for (unsigned i = 0; i < N; ++i) {
+    switch (int alias = IntArr[i]) {
+    default:
+      sideEffect(alias);
+    }
+  }
+  // CHECK: for (auto alias : IntArr)
+  // CHECK-NEXT: switch (alias) {
+
+  for (unsigned i = 0; i < N; ++i) {
+    for (int alias = IntArr[i]; alias < N; ++alias) {
+      sideEffect(alias);
+    }
+  }
+  // CHECK: for (auto alias : IntArr)
+  // CHECK-NEXT: for (; alias < N; ++alias) {
+
+  for (unsigned i = 0; i < N; ++i) {
+    for (unsigned j = 0; int alias = IntArr[i]; ++j) {
+      sideEffect(alias);
+    }
+  }
+  // CHECK: for (auto alias : IntArr)
+  // CHECK-NEXT: for (unsigned j = 0; alias; ++j) {
 }
 
 void refs_and_vals() {