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

llvm-svn: 129300

2 files changed, 23 insertions, 9 deletions

diff --git a/clang/lib/Rewrite/RewriteObjC.cpp b/clang/lib/Rewrite/RewriteObjC.cpp
index 298ed9a9062..ff4c5e91526 100644
--- a/clang/lib/Rewrite/RewriteObjC.cpp
+++ b/clang/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/clang/lib/Rewrite/Rewriter.cpp b/clang/lib/Rewrite/Rewriter.cpp
index c4b6757bacb..3d396d43e4d 100644
--- a/clang/lib/Rewrite/Rewriter.cpp
+++ b/clang/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;
+}