Fixup more objc rwriter bug having to do with

rewriting of blocks which have objective-c
stuff which need be rewritten as well. // rdar://9254348


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

4 files changed, 35 insertions, 12 deletions

diff --git a/include/clang/Rewrite/Rewriter.h b/include/clang/Rewrite/Rewriter.h
index 252f0baf89..397571714c 100644
--- a/include/clang/Rewrite/Rewriter.h
+++ b/include/clang/Rewrite/Rewriter.h
@@ -245,6 +245,10 @@ public:
   /// printer to generate the replacement code.  This returns true if the input
   /// could not be rewritten, or false if successful.
   bool ReplaceStmt(Stmt *From, Stmt *To);
+  
+  /// ConvertToString converts statement 'From' to a string using the
+  /// pretty pronter.
+  std::string ConvertToString(Stmt *From);
 
   /// getEditBuffer - This is like getRewriteBufferFor, but always returns a
   /// buffer, and allows you to write on it directly.  This is useful if you
diff --git a/lib/Rewrite/RewriteObjC.cpp b/lib/Rewrite/RewriteObjC.cpp
index 298ed9a906..ff4c5e9152 100644
--- a/lib/Rewrite/RewriteObjC.cpp
+++ b/lib/Rewrite/RewriteObjC.cpp
@@ -5519,27 +5519,34 @@ Stmt *RewriteObjC::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
   for (Stmt::child_range CI = S->children(); CI; ++CI)
     if (*CI) {
       Stmt *newStmt;
-      Stmt *S = (*CI);
-      if (ObjCIvarRefExpr *IvarRefExpr = dyn_cast<ObjCIvarRefExpr>(S)) {
+      Stmt *ChildStmt = (*CI);
+      if (ObjCIvarRefExpr *IvarRefExpr = dyn_cast<ObjCIvarRefExpr>(ChildStmt)) {
         Expr *OldBase = IvarRefExpr->getBase();
         bool replaced = false;
-        newStmt = RewriteObjCNestedIvarRefExpr(S, replaced);
+        newStmt = RewriteObjCNestedIvarRefExpr(ChildStmt, replaced);
         if (replaced) {
           if (ObjCIvarRefExpr *IRE = dyn_cast<ObjCIvarRefExpr>(newStmt))
             ReplaceStmt(OldBase, IRE->getBase());
           else
-            ReplaceStmt(S, newStmt);
+            ReplaceStmt(ChildStmt, newStmt);
         }
       }
       else
-        newStmt = RewriteFunctionBodyOrGlobalInitializer(S);
-      if (newStmt)
-        *CI = newStmt;
+        newStmt = RewriteFunctionBodyOrGlobalInitializer(ChildStmt);
+      if (newStmt) {
+          if (Expr *PropOrImplicitRefExpr = dyn_cast<Expr>(ChildStmt))
+            if (PropSetters[PropOrImplicitRefExpr] == S) {
+              S = newStmt;
+              newStmt = 0;
+            }
+        if (newStmt)
+          *CI = newStmt;
+      }
       // If dealing with an assignment with LHS being a property reference
       // expression, the entire assignment tree is rewritten into a property
       // setter messaging. This involvs the RHS too. Do not attempt to rewrite
       // RHS again.
-      if (Expr *Exp = dyn_cast<Expr>(S))
+      if (Expr *Exp = dyn_cast<Expr>(ChildStmt))
         if (isa<ObjCPropertyRefExpr>(Exp)) {
           if (PropSetters[Exp]) {
             ++CI;
@@ -5565,7 +5572,7 @@ Stmt *RewriteObjC::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
     PropParentMap = 0;
     ImportedLocalExternalDecls.clear();
     // Now we snarf the rewritten text and stash it away for later use.
-    std::string Str = Rewrite.getRewrittenText(BE->getSourceRange());
+    std::string Str = Rewrite.ConvertToString(BE->getBody());
     RewrittenBlockExprs[BE] = Str;
 
     Stmt *blockTranscribed = SynthBlockInitExpr(BE, InnerBlockDeclRefs);
diff --git a/lib/Rewrite/Rewriter.cpp b/lib/Rewrite/Rewriter.cpp
index c4b6757bac..3d396d43e4 100644
--- a/lib/Rewrite/Rewriter.cpp
+++ b/lib/Rewrite/Rewriter.cpp
@@ -303,3 +303,10 @@ bool Rewriter::ReplaceStmt(Stmt *From, Stmt *To) {
   ReplaceText(From->getLocStart(), Size, Str);
   return false;
 }
+
+std::string Rewriter::ConvertToString(Stmt *From) {
+  std::string SStr;
+  llvm::raw_string_ostream S(SStr);
+  From->printPretty(S, 0, PrintingPolicy(*LangOpts));
+  return SStr;
+}
diff --git a/test/Rewriter/rewrite-block-literal-1.mm b/test/Rewriter/rewrite-block-literal-1.mm
index 681d2ff077..04302d17f1 100644
--- a/test/Rewriter/rewrite-block-literal-1.mm
+++ b/test/Rewriter/rewrite-block-literal-1.mm
@@ -16,9 +16,14 @@ typedef void (^BLOCK_TYPE)(void);
 @implementation CoreDAVTaskGroup
 - (void)_finishInitialSync {
                     CoreDAVTaskGroup *folderPost;
-                    [folderPost setCompletionBlock : (^{
-			self.IVAR = 0;
-                    })];
+  folderPost.completionBlock = ^{
+    self.IVAR = 0;
+    [self _finishInitialSync];
+  };
+
+  [folderPost setCompletionBlock : (^{
+    self.IVAR = 0;
+  })];
 }
 @dynamic IVAR;
 - (void) setCompletionBlock : (BLOCK_TYPE) arg {}