Fix a rewrite bug. // rdar://9039342

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@126435 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 9 insertions, 1 deletions

diff --git a/lib/Rewrite/RewriteObjC.cpp b/lib/Rewrite/RewriteObjC.cpp
index 875a0c7a84..78e370b977 100644
--- a/lib/Rewrite/RewriteObjC.cpp
+++ b/lib/Rewrite/RewriteObjC.cpp
@@ -298,6 +298,7 @@ namespace {
     Stmt *RewriteObjCThrowStmt(ObjCAtThrowStmt *S);
     Stmt *RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S,
                                        SourceLocation OrigEnd);
+    bool IsDeclStmtInForeachHeader(DeclStmt *DS);
     CallExpr *SynthesizeCallToFunctionDecl(FunctionDecl *FD,
                                       Expr **args, unsigned nargs,
                                       SourceLocation StartLoc=SourceLocation(),
@@ -5459,6 +5460,13 @@ Stmt *RewriteObjC::SynthBlockInitExpr(BlockExpr *Exp,
   return NewRep;
 }
 
+bool RewriteObjC::IsDeclStmtInForeachHeader(DeclStmt *DS) {
+  if (const ObjCForCollectionStmt * CS = 
+      dyn_cast<ObjCForCollectionStmt>(Stmts.back()))
+        return CS->getElement() == DS;
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Function Body / Expression rewriting
 //===----------------------------------------------------------------------===//
@@ -5681,7 +5689,7 @@ Stmt *RewriteObjC::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
     //   for (id <FooProtocol> index in someArray) ;
     // This is because RewriteObjCForCollectionStmt() does textual rewriting 
     // and it depends on the original text locations/positions.
-    if (Stmts.empty() || !isa<ObjCForCollectionStmt>(Stmts.back()))
+    if (Stmts.empty() || !IsDeclStmtInForeachHeader(DS))
       RewriteObjCQualifiedInterfaceTypes(*DS->decl_begin());
 
     // Blocks rewrite rules.