RewriteObjC::RewriteObjCForCollectionStmt() needs to handle bodies with a single statement.
Fixes <rdar://problem/6084870> clang ObjC rewriter: for-in enumeration in 1 line produces output with error.
git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@53858 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/Driver/RewriteObjC.cpp b/Driver/RewriteObjC.cpp
index 7bc414c..347874e 100644
--- a/Driver/RewriteObjC.cpp
+++ b/Driver/RewriteObjC.cpp
@@ -1248,9 +1248,25 @@
buf += elementName;
buf += " = nil;\n";
buf += "}\n";
+
// Insert all these *after* the statement body.
- SourceLocation endBodyLoc = OrigEnd.getFileLocWithOffset(1);
- InsertText(endBodyLoc, buf.c_str(), buf.size());
+ if (isa<CompoundStmt>(S->getBody())) {
+ SourceLocation endBodyLoc = OrigEnd.getFileLocWithOffset(1);
+ InsertText(endBodyLoc, buf.c_str(), buf.size());
+ } else {
+ /* Need to treat single statements specially. For example:
+ *
+ * for (A *a in b) if (stuff()) break;
+ * for (A *a in b) xxxyy;
+ *
+ * The following code simply scans ahead to the semi to find the actual end.
+ */
+ const char *stmtBuf = SM->getCharacterData(OrigEnd);
+ const char *semiBuf = strchr(stmtBuf, ';');
+ assert(semiBuf && "Can't find ';'");
+ SourceLocation endBodyLoc = OrigEnd.getFileLocWithOffset(semiBuf-stmtBuf+1);
+ InsertText(endBodyLoc, buf.c_str(), buf.size());
+ }
Stmts.pop_back();
ObjCBcLabelNo.pop_back();
return 0;