Fix a rewriter crash when the whole body of a foreach is itself

rewritten, as in Sema/rewrite-foreach-6.m.  Fariborz/Steve,
please review this to see if it is sane.

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

diff --git a/Driver/RewriteTest.cpp b/Driver/RewriteTest.cpp
index 4c84b8adcb53c330399d87d9435ce3492a19d6bc..d873a3472121963e2bb7b93cbac3372573e310c1 100644 (file)
--- a/Driver/RewriteTest.cpp
+++ b/Driver/RewriteTest.cpp
@@ -203,7 +203,8 @@ namespace {
     Stmt *RewriteObjCCatchStmt(ObjCAtCatchStmt *S);
     Stmt *RewriteObjCFinallyStmt(ObjCAtFinallyStmt *S);
     Stmt *RewriteObjCThrowStmt(ObjCAtThrowStmt *S);
-    Stmt *RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S);
+    Stmt *RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S,
+                                       SourceLocation OrigEnd);
     CallExpr *SynthesizeCallToFunctionDecl(FunctionDecl *FD, 
                                            Expr **args, unsigned nargs);
     Stmt *SynthMessageExpr(ObjCMessageExpr *Exp);
@@ -763,7 +764,9 @@ Stmt *RewriteTest::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
     ObjCBcLabelNo.push_back(++BcLabelCount);
   }
   
-  // Otherwise, just rewrite all children.
+  SourceLocation OrigStmtEnd = S->getLocEnd();
+  
+  // Start by rewriting all children.
   for (Stmt::child_iterator CI = S->child_begin(), E = S->child_end();
        CI != E; ++CI)
     if (*CI) {
@@ -819,7 +822,7 @@ Stmt *RewriteTest::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
   
   if (ObjCForCollectionStmt *StmtForCollection = 
         dyn_cast<ObjCForCollectionStmt>(S))
-    return RewriteObjCForCollectionStmt(StmtForCollection);
+    return RewriteObjCForCollectionStmt(StmtForCollection, OrigStmtEnd);
   if (BreakStmt *StmtBreakStmt =
       dyn_cast<BreakStmt>(S))
     return RewriteBreakStmt(StmtBreakStmt);
@@ -939,7 +942,8 @@ Stmt *RewriteTest::RewriteContinueStmt(ContinueStmt *S) {
 ///       elem = nil;
 ///  }
 ///
-Stmt *RewriteTest::RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S) {
+Stmt *RewriteTest::RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S,
+                                                SourceLocation OrigEnd) {
   assert(!Stmts.empty() && "ObjCForCollectionStmt - Statement stack empty");
   assert(isa<ObjCForCollectionStmt>(Stmts.back()) && 
          "ObjCForCollectionStmt Statement stack mismatch");
@@ -1064,13 +1068,10 @@ Stmt *RewriteTest::RewriteObjCForCollectionStmt(ObjCForCollectionStmt *S) {
   buf += " = nil;\n";
   buf += "}\n";
   // Insert all these *after* the statement body.
-  SourceLocation endBodyLoc = S->getBody()->getLocEnd();
-  const char *endBodyBuf = SM->getCharacterData(endBodyLoc)+1;
-  endBodyLoc = startLoc.getFileLocWithOffset(endBodyBuf-startBuf);
+  SourceLocation endBodyLoc = OrigEnd.getFileLocWithOffset(1);
   Rewrite.InsertText(endBodyLoc, buf.c_str(), buf.size());
   Stmts.pop_back();
   ObjCBcLabelNo.pop_back();
- 
   return 0;
 }
 

diff --git a/test/Sema/rewrite-foreach-6.m b/test/Sema/rewrite-foreach-6.m
new file mode 100644 (file)
index 0000000..bae7e7a
--- /dev/null
+++ b/test/Sema/rewrite-foreach-6.m
@@ -0,0 +1,13 @@
+// RUN: clang %s -rewrite-test
+// rdar://5716356
+// FIXME: Should be able to pipe into clang, but code is not
+// yet correct for other reasons: rdar://5716940
+
+@class NSNotification;
+@class NSMutableArray;
+
+void foo(NSMutableArray *notificationArray, id X) {
+  for (NSNotification *notification in notificationArray)
+    [X postNotification:notification];
+}
+