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

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

4 files changed, 48 insertions, 3 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.
diff --git a/test/Rewriter/rewrite-foreach-5.m b/test/Rewriter/rewrite-foreach-5.m
index adfd7f8378..7baccc3789 100644
--- a/test/Rewriter/rewrite-foreach-5.m
+++ b/test/Rewriter/rewrite-foreach-5.m
@@ -1,4 +1,8 @@
-// RUN: %clang_cc1 -rewrite-objc %s -o -
+// RUN: %clang_cc1 -x objective-c++ -Wno-return-type -fblocks -fms-extensions -rewrite-objc %s -o %t-rw.cpp
+// RUN: %clang_cc1 -fsyntax-only -fblocks -Wno-address-of-temporary -D"id=void*" -D"SEL=void*" -D"__declspec(X)=" %t-rw.cpp
+
+void *sel_registerName(const char *);
+void objc_enumerationMutation(id);
 
 @interface MyList
 - (id) allKeys;
diff --git a/test/Rewriter/rewrite-foreach-6.m b/test/Rewriter/rewrite-foreach-6.m
index 2aa19aecb8..96b472a60e 100644
--- a/test/Rewriter/rewrite-foreach-6.m
+++ b/test/Rewriter/rewrite-foreach-6.m
@@ -1,8 +1,12 @@
-// RUN: %clang_cc1 %s -rewrite-objc -o -
+// RUN: %clang_cc1 -x objective-c++ -Wno-return-type -fblocks -fms-extensions -rewrite-objc %s -o %t-rw.cpp
+// RUN: %clang_cc1 -fsyntax-only -fblocks -Wno-address-of-temporary -D"id=void*" -D"SEL=void*" -D"__declspec(X)=" %t-rw.cpp
 // rdar://5716356
 // FIXME: Should be able to pipe into clang, but code is not
 // yet correct for other reasons: rdar://5716940
 
+void *sel_registerName(const char *);
+void objc_enumerationMutation(id);
+
 @class NSNotification;
 @class NSMutableArray;
 
diff --git a/test/Rewriter/rewrite-foreach-protocol-id.m b/test/Rewriter/rewrite-foreach-protocol-id.m
new file mode 100644
index 0000000000..85d0d0dc44
--- /dev/null
+++ b/test/Rewriter/rewrite-foreach-protocol-id.m
@@ -0,0 +1,29 @@
+// RUN: %clang_cc1 -x objective-c++ -Wno-return-type -fblocks -fms-extensions -rewrite-objc %s -o %t-rw.cpp
+// RUN: %clang_cc1 -fsyntax-only -fblocks -Wno-address-of-temporary -D"id=void*" -D"SEL=void*" -D"__declspec(X)=" %t-rw.cpp
+// rdar:// 9039342
+
+void *sel_registerName(const char *);
+void objc_enumerationMutation(id);
+
+@protocol CoreDAVLeafDataPayload @end
+
+@class NSString;
+
+@interface CoreDAVAction
+- (id) context;
+@end
+
+@interface I
+{
+  id uuidsToAddActions;
+}
+@end
+
+@implementation I
+- (void) Meth {
+  for (id<CoreDAVLeafDataPayload> uuid in uuidsToAddActions) {
+    CoreDAVAction *action = 0;
+    id <CoreDAVLeafDataPayload> payload = [action context];
+  }
+}
+@end