Fix a nested ivar reference rewriting bug.

(Fixes radar 7607605).


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

2 files changed, 63 insertions, 13 deletions

diff --git a/lib/Frontend/RewriteObjC.cpp b/lib/Frontend/RewriteObjC.cpp
index 32cc7472ae..1c4925eeaa 100644
--- a/lib/Frontend/RewriteObjC.cpp
+++ b/lib/Frontend/RewriteObjC.cpp
@@ -278,7 +278,9 @@ namespace {
     ParentMap *PropParentMap; // created lazily.
 
     Stmt *RewriteAtEncode(ObjCEncodeExpr *Exp);
-    Stmt *RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV, SourceLocation OrigStart);
+    Stmt *RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV, SourceLocation OrigStart,
+                                 bool &replaced);
+    Stmt *RewriteObjCNestedIvarRefExpr(Stmt *S, bool &replaced);
     Stmt *RewritePropertyGetter(ObjCPropertyRefExpr *PropRefExpr);
     Stmt *RewritePropertySetter(BinaryOperator *BinOp, Expr *newStmt,
                                 SourceRange SrcRange);
@@ -1209,7 +1211,8 @@ Stmt *RewriteObjC::RewritePropertyGetter(ObjCPropertyRefExpr *PropRefExpr) {
 }
 
 Stmt *RewriteObjC::RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV,
-                                          SourceLocation OrigStart) {
+                                          SourceLocation OrigStart,
+                                          bool &replaced) {
   ObjCIvarDecl *D = IV->getDecl();
   const Expr *BaseExpr = IV->getBase();
   if (CurMethodDef) {
@@ -1238,21 +1241,16 @@ Stmt *RewriteObjC::RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV,
       ParenExpr *PE = new (Context) ParenExpr(IV->getBase()->getLocStart(),
                                                IV->getBase()->getLocEnd(),
                                                castExpr);
+      replaced = true;
       if (IV->isFreeIvar() &&
           CurMethodDef->getClassInterface() == iFaceDecl->getDecl()) {
         MemberExpr *ME = new (Context) MemberExpr(PE, true, D,
                                                    IV->getLocation(),
                                                    D->getType());
-        ReplaceStmt(IV, ME);
         // delete IV; leak for now, see RewritePropertySetter() usage for more info.
         return ME;
       }
-      // Get the old text, only to get its size.
-      std::string SStr;
-      llvm::raw_string_ostream S(SStr);
-      IV->getBase()->printPretty(S, *Context, 0, PrintingPolicy(LangOpts));
       // Get the new text
-      ReplaceStmt(IV->getBase(), PE, S.str().size());
       // Cannot delete IV->getBase(), since PE points to it.
       // Replace the old base with the cast. This is important when doing
       // embedded rewrites. For example, [newInv->_container addObject:0].
@@ -1287,7 +1285,7 @@ Stmt *RewriteObjC::RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV,
       // Don't forget the parens to enforce the proper binding.
       ParenExpr *PE = new (Context) ParenExpr(IV->getBase()->getLocStart(),
                                     IV->getBase()->getLocEnd(), castExpr);
-      ReplaceStmt(IV->getBase(), PE);
+      replaced = true;
       // Cannot delete IV->getBase(), since PE points to it.
       // Replace the old base with the cast. This is important when doing
       // embedded rewrites. For example, [newInv->_container addObject:0].
@@ -1298,6 +1296,24 @@ Stmt *RewriteObjC::RewriteObjCIvarRefExpr(ObjCIvarRefExpr *IV,
   return IV;
 }
 
+Stmt *RewriteObjC::RewriteObjCNestedIvarRefExpr(Stmt *S, bool &replaced) {
+  for (Stmt::child_iterator CI = S->child_begin(), E = S->child_end();
+       CI != E; ++CI) {
+    if (*CI) {
+      Stmt *newStmt = RewriteObjCNestedIvarRefExpr(*CI, replaced);
+      if (newStmt)
+        *CI = newStmt;
+    }
+  }
+  if (ObjCIvarRefExpr *IvarRefExpr = dyn_cast<ObjCIvarRefExpr>(S)) {
+    SourceRange OrigStmtRange = S->getSourceRange();
+    Stmt *newStmt = RewriteObjCIvarRefExpr(IvarRefExpr, OrigStmtRange.getBegin(),
+                                           replaced);
+    return newStmt;
+  }  
+  return S;
+}
+
 /// SynthCountByEnumWithState - To print:
 /// ((unsigned int (*)
 ///  (id, SEL, struct __objcFastEnumerationState *, id *, unsigned int))
@@ -4869,7 +4885,21 @@ Stmt *RewriteObjC::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
   for (Stmt::child_iterator CI = S->child_begin(), E = S->child_end();
        CI != E; ++CI)
     if (*CI) {
-      Stmt *newStmt = RewriteFunctionBodyOrGlobalInitializer(*CI);
+      Stmt *newStmt;
+      Stmt *S = (*CI);
+      if (ObjCIvarRefExpr *IvarRefExpr = dyn_cast<ObjCIvarRefExpr>(S)) {
+        Expr *OldBase = IvarRefExpr->getBase();
+        bool replaced = false;
+        newStmt = RewriteObjCNestedIvarRefExpr(S, replaced);
+        if (replaced) {
+          if (ObjCIvarRefExpr *IRE = dyn_cast<ObjCIvarRefExpr>(newStmt))
+            ReplaceStmt(OldBase, IRE->getBase());
+          else
+            ReplaceStmt(S, newStmt);
+        }
+      }
+      else
+        newStmt = RewriteFunctionBodyOrGlobalInitializer(S);
       if (newStmt)
         *CI = newStmt;
     }
@@ -4891,9 +4921,6 @@ Stmt *RewriteObjC::RewriteFunctionBodyOrGlobalInitializer(Stmt *S) {
   if (ObjCEncodeExpr *AtEncode = dyn_cast<ObjCEncodeExpr>(S))
     return RewriteAtEncode(AtEncode);
 
-  if (ObjCIvarRefExpr *IvarRefExpr = dyn_cast<ObjCIvarRefExpr>(S))
-    return RewriteObjCIvarRefExpr(IvarRefExpr, OrigStmtRange.getBegin());
-
   if (ObjCPropertyRefExpr *PropRefExpr = dyn_cast<ObjCPropertyRefExpr>(S)) {
     BinaryOperator *BinOp = PropSetters[PropRefExpr];
     if (BinOp) {
diff --git a/test/Rewriter/rewrite-cast-ivar-access.mm b/test/Rewriter/rewrite-cast-ivar-access.mm
index 177824d026..c04cbab403 100644
--- a/test/Rewriter/rewrite-cast-ivar-access.mm
+++ b/test/Rewriter/rewrite-cast-ivar-access.mm
@@ -25,6 +25,29 @@ void __CFAssignWithWriteBarrier(void **location, void *value) {
         objc_assign_strongCast((id)value);
 }
 
+// radar 7607605
+@interface RealClass {
+        @public
+        int f;
+}
+@end
+
+@implementation RealClass
+@end
+
+@interface Foo {
+        id reserved;
+}
+@end
+
+@implementation Foo
+- (void)bar {
+        ((RealClass*)reserved)->f = 99;
+}
+@end
+
 // CHECK-LP: ((struct G_IMPL *)arg)->ivar
 
 // CHECK-LP: objc_assign_strongCast((id)value)
+
+// CHECK-LP: ((struct RealClass_IMPL *)((RealClass *)((struct Foo_IMPL *)self)->reserved))->f