[arcmt] Break apart Transforms.cpp.

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

1 files changed, 89 insertions, 1941 deletions

diff --git a/lib/ARCMigrate/Transforms.cpp b/lib/ARCMigrate/Transforms.cpp
index 499c8f034c..0650e3b94f 100644
--- a/lib/ARCMigrate/Transforms.cpp
+++ b/lib/ARCMigrate/Transforms.cpp
@@ -6,151 +6,8 @@
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
-// Transformations:
-//===----------------------------------------------------------------------===//
-//
-// castNonObjCToObjC:
-//
-// A cast of non-objc pointer to an objc one is checked. If the non-objc pointer
-// is from a file-level variable, objc_unretainedObject function is used to
-// convert it.
-//
-//  NSString *str = (NSString *)kUTTypePlainText;
-//  str = b ? kUTTypeRTF : kUTTypePlainText;
-// ---->
-//  NSString *str = objc_unretainedObject(kUTTypePlainText);
-//  str = objc_unretainedObject(b ? kUTTypeRTF : kUTTypePlainText);
-//
-// For a C pointer to ObjC, objc_unretainedPointer is used.
-//
-//  void *vp = str; // NSString*
-// ---->
-//  void *vp = (void*)objc_unretainedPointer(str);
-//
-//===----------------------------------------------------------------------===//
-//
-// rewriteAllocCopyWithZone:
-//
-// Calls to +allocWithZone/-copyWithZone/-mutableCopyWithZone are changed to
-// +alloc/-copy/-mutableCopy if we can safely remove the given parameter.
-//
-//  Foo *foo1 = [[Foo allocWithZone:[self zone]] init];
-// ---->
-//  Foo *foo1 = [[Foo alloc] init];
-//
-//===----------------------------------------------------------------------===//
-//
-// rewriteAutoreleasePool:
-//
-// Calls to NSAutoreleasePools will be rewritten as an @autorelease scope.
-//
-//  NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
-//  ...
-//  [pool release];
-// ---->
-//  @autorelease {
-//  ...
-//  }
-//
-// An NSAutoreleasePool will not be touched if:
-// - There is not a corresponding -release/-drain in the same scope
-// - Not all references of the NSAutoreleasePool variable can be removed
-// - There is a variable that is declared inside the intended @autorelease scope
-//   which is also used outside it.
-//
-//===----------------------------------------------------------------------===//
-//
-// makeAssignARCSafe:
-//
-// Add '__strong' where appropriate.
-//
-//  for (id x in collection) {
-//    x = 0;
-//  }
-// ---->
-//  for (__strong id x in collection) {
-//    x = 0;
-//  }
-//
-//===----------------------------------------------------------------------===//
-//
-// removeRetainReleaseDealloc:
-//
-// Removes retain/release/autorelease/dealloc messages.
-//
-//  return [[foo retain] autorelease];
-// ---->
-//  return foo;
-//
-//===----------------------------------------------------------------------===//
-//
-// removeEmptyStatements:
-//
-// Removes empty statements that are leftovers from previous transformations.
-// e.g for
-//
-//  [x retain];
-//
-// removeRetainReleaseDealloc will leave an empty ";" that removeEmptyStatements
-// will remove.
-//
-//===----------------------------------------------------------------------===//
-//
-// changeIvarsOfAssignProperties:
-//
-// If a property is synthesized with 'assign' attribute and the user didn't
-// set a lifetime attribute, change the property to 'weak' or add
-// __unsafe_unretained if the ARC runtime is not available.
-//
-//  @interface Foo : NSObject {
-//      NSObject *x;
-//  }
-//  @property (assign) id x;
-//  @end
-// ---->
-//  @interface Foo : NSObject {
-//      NSObject *__weak x;
-//  }
-//  @property (weak) id x;
-//  @end
-//
-//===----------------------------------------------------------------------===//
-//
-// rewriteUnusedDelegateInit:
-//
-// Rewrites an unused result of calling a delegate initialization, to assigning
-// the result to self.
-// e.g
-//  [self init];
-// ---->
-//  self = [self init];
-//
-//===----------------------------------------------------------------------===//
-//
-// rewriteBlockObjCVariable:
-//
-// Adding __block to an obj-c variable could be either because the the variable
-// is used for output storage or the user wanted to break a retain cycle.
-// This transformation checks whether a reference of the variable for the block
-// is actually needed (it is assigned to or its address is taken) or not.
-// If the reference is not needed it will assume __block was added to break a
-// cycle so it will remove '__block' and add __weak/__unsafe_unretained.
-// e.g
-//
-//   __block Foo *x;
-//   bar(^ { [x cake]; });
-// ---->
-//   __weak Foo *x;
-//   bar(^ { [x cake]; });
-//
-//===----------------------------------------------------------------------===//
-//
-// removeZeroOutIvarsInDealloc:
-//
-// Removes zero'ing out "strong" @synthesized properties in a -dealloc method.
-//
-//===----------------------------------------------------------------------===//
 
+#include "Transforms.h"
 #include "Internals.h"
 #include "clang/Sema/SemaDiagnostic.h"
 #include "clang/AST/RecursiveASTVisitor.h"
@@ -165,82 +22,51 @@
 
 using namespace clang;
 using namespace arcmt;
+using namespace trans;
 using llvm::StringRef;
 
 //===----------------------------------------------------------------------===//
-// Transformations.
+// Helpers.
 //===----------------------------------------------------------------------===//
 
-namespace {
+/// \brief 'Loc' is the end of a statement range. This returns the location
+/// immediately after the semicolon following the statement.
+/// If no semicolon is found or the location is inside a macro, the returned
+/// source location will be invalid.
+SourceLocation trans::findLocationAfterSemi(SourceLocation loc,
+                                            ASTContext &Ctx) {
+  SourceManager &SM = Ctx.getSourceManager();
+  if (loc.isMacroID()) {
+    if (!SM.isAtEndOfMacroInstantiation(loc))
+      return SourceLocation();
+    loc = SM.getInstantiationRange(loc).second;
+  }
+  loc = Lexer::getLocForEndOfToken(loc, /*Offset=*/0, SM, Ctx.getLangOptions());
 
-class RemovablesCollector : public RecursiveASTVisitor<RemovablesCollector> {
-  llvm::DenseSet<Expr *> &Removables;
+  // Break down the source location.
+  std::pair<FileID, unsigned> locInfo = SM.getDecomposedLoc(loc);
 
-public:
-  RemovablesCollector(llvm::DenseSet<Expr *> &removables)
-  : Removables(removables) { }
-  
-  bool shouldWalkTypesOfTypeLocs() const { return false; }
-  
-  bool TraverseStmtExpr(StmtExpr *E) {
-    CompoundStmt *S = E->getSubStmt();
-    for (CompoundStmt::body_iterator
-        I = S->body_begin(), E = S->body_end(); I != E; ++I) {
-      if (I != E - 1)
-        mark(*I);
-      TraverseStmt(*I);
-    }
-    return true;
-  }
-  
-  bool VisitCompoundStmt(CompoundStmt *S) {
-    for (CompoundStmt::body_iterator
-        I = S->body_begin(), E = S->body_end(); I != E; ++I)
-      mark(*I);
-    return true;
-  }
-  
-  bool VisitIfStmt(IfStmt *S) {
-    mark(S->getThen());
-    mark(S->getElse());
-    return true;
-  }
-  
-  bool VisitWhileStmt(WhileStmt *S) {
-    mark(S->getBody());
-    return true;
-  }
-  
-  bool VisitDoStmt(DoStmt *S) {
-    mark(S->getBody());
-    return true;
-  }
-  
-  bool VisitForStmt(ForStmt *S) {
-    mark(S->getInit());
-    mark(S->getInc());
-    mark(S->getBody());
-    return true;
-  }
-  
-private:
-  void mark(Stmt *S) {
-    if (!S) return;
-    
-    if (LabelStmt *Label = dyn_cast<LabelStmt>(S))
-      return mark(Label->getSubStmt());
-    if (ImplicitCastExpr *CE = dyn_cast<ImplicitCastExpr>(S))
-      return mark(CE->getSubExpr());
-    if (ExprWithCleanups *EWC = dyn_cast<ExprWithCleanups>(S))
-      return mark(EWC->getSubExpr());
-    if (Expr *E = dyn_cast<Expr>(S))
-      Removables.insert(E);
-  }
-};
+  // Try to load the file buffer.
+  bool invalidTemp = false;
+  llvm::StringRef file = SM.getBufferData(locInfo.first, &invalidTemp);
+  if (invalidTemp)
+    return SourceLocation();
+
+  const char *tokenBegin = file.data() + locInfo.second;
+
+  // Lex from the start of the given location.
+  Lexer lexer(SM.getLocForStartOfFile(locInfo.first),
+              Ctx.getLangOptions(),
+              file.begin(), tokenBegin, file.end());
+  Token tok;
+  lexer.LexFromRawLexer(tok);
+  if (tok.isNot(tok::semi))
+    return SourceLocation();
 
-} // end anonymous namespace.
+  return tok.getLocation().getFileLocWithOffset(1);
+}
 
-static bool HasSideEffects(Expr *E, ASTContext &Ctx) {
+bool trans::hasSideEffects(Expr *E, ASTContext &Ctx) {
   if (!E || !E->HasSideEffects(Ctx))
     return false;
 
@@ -257,7 +83,7 @@ static bool HasSideEffects(Expr *E, ASTContext &Ctx) {
     case ObjCMessageExpr::SuperInstance:
       return false;
     case ObjCMessageExpr::Instance:
-      return HasSideEffects(ME->getInstanceReceiver(), Ctx);
+      return hasSideEffects(ME->getInstanceReceiver(), Ctx);
     default:
       break;
     }
@@ -269,66 +95,23 @@ static bool HasSideEffects(Expr *E, ASTContext &Ctx) {
   return true;
 }
 
-static void removeDeallocMethod(MigrationPass &pass) {
-    ASTContext &Ctx = pass.Ctx;
-    TransformActions &TA = pass.TA;
-    DeclContext *DC = Ctx.getTranslationUnitDecl();
-    ObjCMethodDecl *DeallocMethodDecl = 0;
-    IdentifierInfo *II = &Ctx.Idents.get("dealloc");
-    
-    for (DeclContext::decl_iterator I = DC->decls_begin(), E = DC->decls_end();
-         I != E; ++I) {
-        Decl *D = *I;
-        if (ObjCImplementationDecl *IMD = 
-            dyn_cast<ObjCImplementationDecl>(D)) {
-            DeallocMethodDecl = 0;
-            for (ObjCImplementationDecl::instmeth_iterator I = 
-                 IMD->instmeth_begin(), E = IMD->instmeth_end();
-                 I != E; ++I) {
-                ObjCMethodDecl *OMD = *I;
-                if (OMD->isInstanceMethod() &&
-                    OMD->getSelector() == Ctx.Selectors.getSelector(0, &II)) {
-                    DeallocMethodDecl = OMD;
-                    break;
-                }
-            }
-            if (DeallocMethodDecl && 
-                DeallocMethodDecl->getCompoundBody()->body_empty()) {
-              Transaction Trans(TA);
-              TA.remove(DeallocMethodDecl->getSourceRange());
-            }
-        }
-    }
-}
-
 namespace {
 
 class ReferenceClear : public RecursiveASTVisitor<ReferenceClear> {
-  llvm::DenseSet<Expr *> &Refs;
+  ExprSet &Refs;
 public:
-  ReferenceClear(llvm::DenseSet<Expr *> &refs) : Refs(refs) { }
+  ReferenceClear(ExprSet &refs) : Refs(refs) { }
   bool VisitDeclRefExpr(DeclRefExpr *E) { Refs.erase(E); return true; }
   bool VisitBlockDeclRefExpr(BlockDeclRefExpr *E) { Refs.erase(E); return true; }
-  void clearRefsIn(Stmt *S) { TraverseStmt(S); }
-  template <typename iterator>
-  void clearRefsIn(iterator begin, iterator end) {
-    for (; begin != end; ++begin)
-      TraverseStmt(*begin);
-  }
 };
 
 class ReferenceCollector : public RecursiveASTVisitor<ReferenceCollector> {
   ValueDecl *Dcl;
-  llvm::DenseSet<Expr *> &Refs;
+  ExprSet &Refs;
 
 public:
-  ReferenceCollector(llvm::DenseSet<Expr *> &refs)
-    : Dcl(0), Refs(refs) { }
-
-  void lookFor(ValueDecl *D, Stmt *S) {
-    Dcl = D;
-    TraverseStmt(S);
-  }
+  ReferenceCollector(ValueDecl *D, ExprSet &refs)
+    : Dcl(D), Refs(refs) { }
 
   bool VisitDeclRefExpr(DeclRefExpr *E) {
     if (E->getDecl() == Dcl)
@@ -343,1718 +126,83 @@ public:
   }
 };
 
-class ReleaseCollector : public RecursiveASTVisitor<ReleaseCollector> {
-  Decl *Dcl;
-  llvm::SmallVectorImpl<ObjCMessageExpr *> &Releases;
-
-public:
-  ReleaseCollector(Decl *D, llvm::SmallVectorImpl<ObjCMessageExpr *> &releases)
-    : Dcl(D), Releases(releases) { }
-
-  bool VisitObjCMessageExpr(ObjCMessageExpr *E) {
-    if (!E->isInstanceMessage())
-      return true;
-    if (E->getMethodFamily() != OMF_release)
-      return true;
-    Expr *instance = E->getInstanceReceiver()->IgnoreParenCasts();
-    if (DeclRefExpr *DE = dyn_cast<DeclRefExpr>(instance)) {
-      if (DE->getDecl() == Dcl)
-        Releases.push_back(E);
-    }
-    return true;
-  }
-};
-
-template <typename BODY_TRANS>
-class BodyTransform : public RecursiveASTVisitor<BodyTransform<BODY_TRANS> > {
-  MigrationPass &Pass;
-
-public:
-  BodyTransform(MigrationPass &pass) : Pass(pass) { }
-
-  void handleBody(Decl *D) {
-    Stmt *body = D->getBody();
-    if (body) {
-      BODY_TRANS(D, Pass).transformBody(body);
-    }
-  }
-
-  bool TraverseBlockDecl(BlockDecl *D) {
-    handleBody(D);
-    return true;
-  }
-  bool TraverseObjCMethodDecl(ObjCMethodDecl *D) {
-    if (D->isThisDeclarationADefinition())
-      handleBody(D);
-    return true;
-  }
-  bool TraverseFunctionDecl(FunctionDecl *D) {
-    if (D->isThisDeclarationADefinition())
-      handleBody(D);
-    return true;
-  }
-};
-
-} // anonymous namespace
-
-//===----------------------------------------------------------------------===//
-// makeAssignARCSafe
-//===----------------------------------------------------------------------===//
-
-namespace {
-
-class ARCAssignChecker : public RecursiveASTVisitor<ARCAssignChecker> {
-  MigrationPass &Pass;
-  llvm::DenseSet<VarDecl *> ModifiedVars;
-
-public:
-  ARCAssignChecker(MigrationPass &pass) : Pass(pass) { }
-
-  bool VisitBinaryOperator(BinaryOperator *Exp) {
-    Expr *E = Exp->getLHS();
-    SourceLocation OrigLoc = E->getExprLoc();
-    SourceLocation Loc = OrigLoc;
-    DeclRefExpr *declRef = dyn_cast<DeclRefExpr>(E->IgnoreParenCasts());
-    if (declRef && isa<VarDecl>(declRef->getDecl())) {
-      ASTContext &Ctx = Pass.Ctx;
-      Expr::isModifiableLvalueResult IsLV = E->isModifiableLvalue(Ctx, &Loc);
-      if (IsLV != Expr::MLV_ConstQualified)
-        return true;
-      VarDecl *var = cast<VarDecl>(declRef->getDecl());
-      if (var->isARCPseudoStrong()) {
-        Transaction Trans(Pass.TA);
-        if (Pass.TA.clearDiagnostic(diag::err_typecheck_arr_assign_enumeration,
-                                    Exp->getOperatorLoc())) {
-          if (!ModifiedVars.count(var)) {
-            TypeLoc TLoc = var->getTypeSourceInfo()->getTypeLoc();
-            Pass.TA.insert(TLoc.getBeginLoc(), "__strong ");
-            ModifiedVars.insert(var);
-          }
-        }
-      }
-    }
-    
-    return true;
-  }
-};
-
-} // anonymous namespace
-
-static void makeAssignARCSafe(MigrationPass &pass) {
-  ARCAssignChecker assignCheck(pass);
-  assignCheck.TraverseDecl(pass.Ctx.getTranslationUnitDecl());
-}
-
-//===----------------------------------------------------------------------===//
-// castNonObjCToObjC
-//===----------------------------------------------------------------------===//
-
-namespace {
-
-class NonObjCToObjCCaster : public RecursiveASTVisitor<NonObjCToObjCCaster> {
-  MigrationPass &Pass;
-  IdentifierInfo *SelfII;
-public:
-  NonObjCToObjCCaster(MigrationPass &pass) : Pass(pass) {
-    SelfII = &Pass.Ctx.Idents.get("self");
-  }
-
-  bool VisitCastExpr(CastExpr *E) {
-    if (E->getCastKind() != CK_AnyPointerToObjCPointerCast
-        && E->getCastKind() != CK_BitCast)
-      return true;
-
-    QualType castType = E->getType();
-    Expr *castExpr = E->getSubExpr();
-    QualType castExprType = castExpr->getType();
-
-    if (castType->isObjCObjectPointerType() &&
-        castExprType->isObjCObjectPointerType())
-      return true;
-    if (!castType->isObjCObjectPointerType() &&
-        !castExprType->isObjCObjectPointerType())
-      return true;
-    
-    bool exprRetainable = castExprType->isObjCIndirectLifetimeType();
-    bool castRetainable = castType->isObjCIndirectLifetimeType();
-    if (exprRetainable == castRetainable) return true;
-
-    if (castExpr->isNullPointerConstant(Pass.Ctx,
-                                        Expr::NPC_ValueDependentIsNull))
-      return true;
-
-    SourceLocation loc = castExpr->getExprLoc();
-    if (loc.isValid() && Pass.Ctx.getSourceManager().isInSystemHeader(loc))
-      return true;
-
-    if (castType->isObjCObjectPointerType())
-      transformNonObjCToObjCCast(E);
-    else
-      transformObjCToNonObjCCast(E);
-
-    return true;
-  }
-
-private:
-  void transformNonObjCToObjCCast(CastExpr *E) {
-    if (!E) return;
-
-    // Global vars are assumed that are cast as unretained.
-    if (isGlobalVar(E))
-      if (E->getSubExpr()->getType()->isPointerType()) {
-        castToObjCObject(E, /*retained=*/false);
-        return;
-      }
-
-    // If the cast is directly over the result of a Core Foundation function
-    // try to figure out whether it should be cast as retained or unretained.
-    Expr *inner = E->IgnoreParenCasts();
-    if (CallExpr *callE = dyn_cast<CallExpr>(inner)) {
-      if (FunctionDecl *FD = callE->getDirectCallee()) {
-        if (FD->getAttr<CFReturnsRetainedAttr>()) {
-          castToObjCObject(E, /*retained=*/true);
-          return;
-        }
-        if (FD->getAttr<CFReturnsNotRetainedAttr>()) {
-          castToObjCObject(E, /*retained=*/false);
-          return;
-        }
-        if (FD->isGlobal() &&
-            FD->getIdentifier() &&
-            ento::cocoa::isRefType(E->getSubExpr()->getType(), "CF",
-                                   FD->getIdentifier()->getName())) {
-          StringRef fname = FD->getIdentifier()->getName();
-          if (fname.endswith("Retain") ||
-              fname.find("Create") != StringRef::npos ||
-              fname.find("Copy") != StringRef::npos) {
-            castToObjCObject(E, /*retained=*/true);
-            return;
-          }
-
-          if (fname.find("Get") != StringRef::npos) {
-            castToObjCObject(E, /*retained=*/false);
-            return;
-          }
-        }
-      }
-    }
-  }
-
-  void castToObjCObject(CastExpr *E, bool retained) {
-    rewriteToBridgedCast(E, retained ? OBC_BridgeTransfer : OBC_Bridge);
-  }
-
-  void rewriteToBridgedCast(CastExpr *E, ObjCBridgeCastKind Kind) {
-    TransformActions &TA = Pass.TA;
-
-    // We will remove the compiler diagnostic.
-    if (!TA.hasDiagnostic(diag::err_arc_mismatched_cast,
-                          diag::err_arc_cast_requires_bridge,
-                          E->getLocStart()))
-      return;
-
-    StringRef bridge;
-    switch(Kind) {
-    case OBC_Bridge:
-      bridge = "__bridge "; break;
-    case OBC_BridgeTransfer:
-      bridge = "__bridge_transfer "; break;
-    case OBC_BridgeRetained:
-      bridge = "__bridge_retained "; break;
-    }
-
-    Transaction Trans(TA);
-    TA.clearDiagnostic(diag::err_arc_mismatched_cast,
-                       diag::err_arc_cast_requires_bridge,
-                       E->getLocStart());
-    if (CStyleCastExpr *CCE = dyn_cast<CStyleCastExpr>(E)) {
-      TA.insertAfterToken(CCE->getLParenLoc(), bridge);
-    } else {
-      SourceLocation insertLoc = E->getSubExpr()->getLocStart();
-      llvm::SmallString<128> newCast;
-      newCast += '(';
-      newCast += bridge;
-      newCast += E->getType().getAsString(Pass.Ctx.PrintingPolicy);
-      newCast += ')';
-
-      if (isa<ParenExpr>(E->getSubExpr())) {
-        TA.insert(insertLoc, newCast.str());
-      } else {
-        newCast += '(';
-        TA.insert(insertLoc, newCast.str());
-        TA.insertAfterToken(E->getLocEnd(), ")");
-      }
-    }
-  }
-
-  void transformObjCToNonObjCCast(CastExpr *E) {
-    if (isSelf(E->getSubExpr()))
-      return rewriteToBridgedCast(E, OBC_Bridge);
-  }
-
-  bool isSelf(Expr *E) {
-    E = E->IgnoreParenLValueCasts();
-    if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(E))
-      if (DRE->getDecl()->getIdentifier() == SelfII)
-        return true;
-    return false;
-  }
-
-  static bool isGlobalVar(Expr *E) {
-    E = E->IgnoreParenCasts();
-    if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(E))
-      return DRE->getDecl()->getDeclContext()->isFileContext();
-    if (ConditionalOperator *condOp = dyn_cast<ConditionalOperator>(E))
-      return isGlobalVar(condOp->getTrueExpr()) &&
-             isGlobalVar(condOp->getFalseExpr());
-
-    return false;  
-  }
-};
-
-} // end anonymous namespace
-
-static void castNonObjCToObjC(MigrationPass &pass) {
-  NonObjCToObjCCaster trans(pass);
-  trans.TraverseDecl(pass.Ctx.getTranslationUnitDecl());
-}
-
-//===----------------------------------------------------------------------===//
-// rewriteAllocCopyWithZone
-//===----------------------------------------------------------------------===//
-
-namespace {
-
-class AllocCopyWithZoneRewriter :
-                         public RecursiveASTVisitor<AllocCopyWithZoneRewriter> {
-  Decl *Dcl;
-  Stmt *Body;
-  MigrationPass &Pass;
-
-  Selector allocWithZoneSel;
-  Selector copyWithZoneSel;
-  Selector mutableCopyWithZoneSel;
-  Selector zoneSel;
-  IdentifierInfo *NSZoneII;
-
-  std::vector<DeclStmt *> NSZoneVars;
-  std::vector<Expr *> Removals;
-
-public:
-  AllocCopyWithZoneRewriter(Decl *D, MigrationPass &pass)
-    : Dcl(D), Body(0), Pass(pass) {
-    SelectorTable &sels = pass.Ctx.Selectors;
-    IdentifierTable &ids = pass.Ctx.Idents; 
-    allocWithZoneSel = sels.getUnarySelector(&ids.get("allocWithZone"));
-    copyWithZoneSel = sels.getUnarySelector(&ids.get("copyWithZone"));
-    mutableCopyWithZoneSel = sels.getUnarySelector(
-                                               &ids.get("mutableCopyWithZone"));
-    zoneSel = sels.getNullarySelector(&ids.get("zone"));
-    NSZoneII = &ids.get("_NSZone");
-  }
-
-  void transformBody(Stmt *body) {
-    Body = body;
-    // Don't change allocWithZone/copyWithZone messages inside
-    // custom implementations of such methods, it can lead to infinite loops.
-    if (ObjCMethodDecl *MD = dyn_cast<ObjCMethodDecl>(Dcl)) {
-      Selector sel = MD->getSelector();
-      if (sel == allocWithZoneSel ||
-          sel == copyWithZoneSel ||
-          sel == mutableCopyWithZoneSel ||
-          sel == zoneSel)
-        return;
-    }
-
-    TraverseStmt(body);
-  }
-
-  ~AllocCopyWithZoneRewriter() {
-    for (std::vector<DeclStmt *>::reverse_iterator
-           I = NSZoneVars.rbegin(), E = NSZoneVars.rend(); I != E; ++I) {
-      DeclStmt *DS = *I;
-      DeclGroupRef group = DS->getDeclGroup();
-      std::vector<Expr *> varRemovals = Removals;
-
-      bool areAllVarsUnused = true;
-      for (std::reverse_iterator<DeclGroupRef::iterator>
-             DI(group.end()), DE(group.begin()); DI != DE; ++DI) {
-        VarDecl *VD = cast<VarDecl>(*DI);
-        if (isNSZoneVarUsed(VD, varRemovals)) {
-          areAllVarsUnused = false;
-          break;
-        }
-        varRemovals.push_back(VD->getInit());
-      }
-
-      if (areAllVarsUnused) {
-        Transaction Trans(Pass.TA);
-        clearUnavailableDiags(DS);
-        Pass.TA.removeStmt(DS);
-        Removals.swap(varRemovals);
-      }
-    }
-  }
-
-  bool VisitObjCMessageExpr(ObjCMessageExpr *E) {
-    if (!isAllocCopyWithZoneCall(E))
-      return true;
-    Expr *arg = E->getArg(0);
-    if (paramToAllocWithZoneHasSideEffects(arg))
-      return true;
-
-    Pass.TA.startTransaction();
-
-    clearUnavailableDiags(arg);
-    Pass.TA.clearDiagnostic(diag::err_unavailable_message,
-                            E->getReceiverRange().getBegin());
-
-    Pass.TA.remove(SourceRange(E->getSelectorLoc(), arg->getLocEnd()));
-    StringRef rewrite;
-    if (E->getSelector() == allocWithZoneSel)
-      rewrite = "alloc";
-    else if (E->getSelector() == copyWithZoneSel)
-      rewrite = "copy";
-    else {
-      assert(E->getSelector() == mutableCopyWithZoneSel);
-      rewrite = "mutableCopy";
-    }
-    Pass.TA.insert(E->getSelectorLoc(), rewrite);
-
-    bool failed = Pass.TA.commitTransaction();
-    if (!failed)
-      Removals.push_back(arg);
-
-    return true;
-  }
-
-  bool VisitDeclStmt(DeclStmt *DS) {
-    DeclGroupRef group = DS->getDeclGroup();
-    if (group.begin() == group.end())
-      return true;
-    for (DeclGroupRef::iterator
-           DI = group.begin(), DE = group.end(); DI != DE; ++DI)
-      if (!isRemovableNSZoneVar(*DI))
-        return true;
-
-    NSZoneVars.push_back(DS);
-    return true;
-  }
-
-private:
-  bool isRemovableNSZoneVar(Decl *D) {
-    if (VarDecl *VD = dyn_cast<VarDecl>(D)) {
-      if (isNSZone(VD->getType()))
-        return !paramToAllocWithZoneHasSideEffects(VD->getInit());
-    }
-    return false;
-  }
-
-  bool isNSZone(RecordDecl *RD) {
-    return RD && RD->getIdentifier() == NSZoneII;
-  }
-
-  bool isNSZone(QualType Ty) {
-    QualType pointee = Ty->getPointeeType();
-    if (pointee.isNull())
-      return false;
-    if (const RecordType *recT = pointee->getAsStructureType())
-      return isNSZone(recT->getDecl());
-    return false;
-  }
-
-  bool isNSZoneVarUsed(VarDecl *D, std::vector<Expr *> &removals) {
-    llvm::DenseSet<Expr *> refs;
-
-    ReferenceCollector refColl(refs);
-    refColl.lookFor(D, Body);
-
-    ReferenceClear refClear(refs);
-    refClear.clearRefsIn(removals.begin(), removals.end());
-
-    return !refs.empty();
-  }
-
-  bool isAllocCopyWithZoneCall(ObjCMessageExpr *E) {
-    if (E->getNumArgs() == 1 &&
-        E->getSelector() == allocWithZoneSel &&
-        (E->isClassMessage() ||
-         Pass.TA.hasDiagnostic(diag::err_unavailable_message,
-                               E->getReceiverRange().getBegin())))
-      return true;
-
-    return E->isInstanceMessage() &&
-           E->getNumArgs() == 1   &&
-           (E->getSelector() == copyWithZoneSel ||
-            E->getSelector() == mutableCopyWithZoneSel);
-  }
-
-  bool isZoneCall(ObjCMessageExpr *E) {
-    return E->isInstanceMessage() &&
-           E->getNumArgs() == 0   &&
-           E->getSelector() == zoneSel;
-  }
-
-  bool paramToAllocWithZoneHasSideEffects(Expr *E) {
-    if (!HasSideEffects(E, Pass.Ctx))
-      return false;
-    E = E->IgnoreParenCasts();
-    ObjCMessageExpr *ME = dyn_cast<ObjCMessageExpr>(E);
-    if (!ME)
-      return true;
-    if (!isZoneCall(ME))
-      return true;
-    return HasSideEffects(ME->getInstanceReceiver(), Pass.Ctx);
-  }
-
-  void clearUnavailableDiags(Stmt *S) {
-    if (S)
-      Pass.TA.clearDiagnostic(diag::err_unavailable,
-                              diag::err_unavailable_message,
-                              S->getSourceRange());
-  }
-};
-
-} // end anonymous namespace
-
-static void rewriteAllocCopyWithZone(MigrationPass &pass) {
-  BodyTransform<AllocCopyWithZoneRewriter> trans(pass);
-  trans.TraverseDecl(pass.Ctx.getTranslationUnitDecl());
-}
-
-//===----------------------------------------------------------------------===//
-// rewriteAutoreleasePool
-//===----------------------------------------------------------------------===//
-
-/// \brief 'Loc' is the end of a statement range. This returns the location
-/// immediately after the semicolon following the statement.
-/// If no semicolon is found or the location is inside a macro, the returned
-/// source location will be invalid.
-static SourceLocation findLocationAfterSemi(ASTContext &Ctx,
-                                            SourceLocation loc) {
-  SourceManager &SM = Ctx.getSourceManager();
-  if (loc.isMacroID()) {
-    if (!SM.isAtEndOfMacroInstantiation(loc))
-      return SourceLocation();
-    loc = SM.getInstantiationRange(loc).second;
-  }
-  loc = Lexer::getLocForEndOfToken(loc, /*Offset=*/0, SM, Ctx.getLangOptions());
-
-  // Break down the source location.
-  std::pair<FileID, unsigned> locInfo = SM.getDecomposedLoc(loc);
-
-  // Try to load the file buffer.
-  bool invalidTemp = false;
-  llvm::StringRef file = SM.getBufferData(locInfo.first, &invalidTemp);
-  if (invalidTemp)
-    return SourceLocation();
-
-  const char *tokenBegin = file.data() + locInfo.second;
-
-  // Lex from the start of the given location.
-  Lexer lexer(SM.getLocForStartOfFile(locInfo.first),
-              Ctx.getLangOptions(),
-              file.begin(), tokenBegin, file.end());
-  Token tok;
-  lexer.LexFromRawLexer(tok);
-  if (tok.isNot(tok::semi))
-    return SourceLocation();
-
-  return tok.getLocation().getFileLocWithOffset(1);
-}
-
-namespace {
+class RemovablesCollector : public RecursiveASTVisitor<RemovablesCollector> {
+  ExprSet &Removables;
 
-class AutoreleasePoolRewriter
-                         : public RecursiveASTVisitor<AutoreleasePoolRewriter> {
 public:
-  AutoreleasePoolRewriter(Decl *D, MigrationPass &pass)
-    : Dcl(D), Body(0), Pass(pass) {
-    PoolII = &pass.Ctx.Idents.get("NSAutoreleasePool");
-    DrainSel = pass.Ctx.Selectors.getNullarySelector(
-                                                 &pass.Ctx.Idents.get("drain"));
-  }
-
-  void transformBody(Stmt *body) {
-    Body = body;
-    TraverseStmt(body);
-  }
-  
-  ~AutoreleasePoolRewriter() {
-    llvm::SmallVector<VarDecl *, 8> VarsToHandle;
-
-    for (std::map<VarDecl *, PoolVarInfo>::iterator
-           I = PoolVars.begin(), E = PoolVars.end(); I != E; ++I) {
-      VarDecl *var = I->first;
-      PoolVarInfo &info = I->second;
-
-      // Check that we can handle/rewrite all references of the pool.
-
-      ReferenceClear refClear(info.Refs);
-      refClear.clearRefsIn(info.Dcl);
-      for (llvm::SmallVectorImpl<PoolScope>::iterator
-             scpI = info.Scopes.begin(),
-             scpE = info.Scopes.end(); scpI != scpE; ++scpI) {
-        PoolScope &scope = *scpI;
-        refClear.clearRefsIn(*scope.Begin);
-        refClear.clearRefsIn(*scope.End);
-        refClear.clearRefsIn(scope.Releases.begin(), scope.Releases.end());
-      }
-
-      // Even if one reference is not handled we will not do anything about that
-      // pool variable.
-      if (info.Refs.empty())
-        VarsToHandle.push_back(var);
-    }
-
-    for (unsigned i = 0, e = VarsToHandle.size(); i != e; ++i) {
-      PoolVarInfo &info = PoolVars[VarsToHandle[i]];
-
-      Transaction Trans(Pass.TA);
-
-      clearUnavailableDiags(info.Dcl);
-      Pass.TA.removeStmt(info.Dcl);
-
-      // Add "@autoreleasepool { }"
-      for (llvm::SmallVectorImpl<PoolScope>::iterator
-             scpI = info.Scopes.begin(),
-             scpE = info.Scopes.end(); scpI != scpE; ++scpI) {
-        PoolScope &scope = *scpI;
-        clearUnavailableDiags(*scope.Begin);
-        clearUnavailableDiags(*scope.End);
-        if (scope.IsFollowedBySimpleReturnStmt) {
-          // Include the return in the scope.
-          Pass.TA.replaceStmt(*scope.Begin, "@autoreleasepool {");
-          Pass.TA.removeStmt(*scope.End);
-          Stmt::child_iterator retI = scope.End;
-          ++retI;
-          SourceLocation afterSemi = findLocationAfterSemi(Pass.Ctx,
-                                                          (*retI)->getLocEnd());
-          assert(afterSemi.isValid() &&
-                 "Didn't we check before setting IsFollowedBySimpleReturnStmt "
-                 "to true?");
-          Pass.TA.insertAfterToken(afterSemi, "\n}");
-          Pass.TA.increaseIndentation(
-                                SourceRange(scope.getIndentedRange().getBegin(),
-                                            (*retI)->getLocEnd()),
-                                      scope.CompoundParent->getLocStart());
-        } else {
-          Pass.TA.replaceStmt(*scope.Begin, "@autoreleasepool {");
-          Pass.TA.replaceStmt(*scope.End, "}");
-          Pass.TA.increaseIndentation(scope.getIndentedRange(),
-                                      scope.CompoundParent->getLocStart());
-        }
-      }
-
-      // Remove rest of pool var references.
-      for (llvm::SmallVectorImpl<PoolScope>::iterator
-             scpI = info.Scopes.begin(),
-             scpE = info.Scopes.end(); scpI != scpE; ++scpI) {
-        PoolScope &scope = *scpI;
-        for (llvm::SmallVectorImpl<ObjCMessageExpr *>::iterator
-               relI = scope.Releases.begin(),
-               relE = scope.Releases.end(); relI != relE; ++relI) {
-          clearUnavailableDiags(*relI);
-          Pass.TA.removeStmt(*relI);
-        }
-      }
-    }
-  }
-
-  bool VisitCompoundStmt(CompoundStmt *S) {
-    llvm::SmallVector<PoolScope, 4> Scopes;
-
-    for (Stmt::child_iterator
-           I = S->body_begin(), E = S->body_end(); I != E; ++I) {
-      Stmt *child = getEssential(*I);
-      if (DeclStmt *DclS = dyn_cast<DeclStmt>(child)) {
-        if (DclS->isSingleDecl()) {
-          if (VarDecl *VD = dyn_cast<VarDecl>(DclS->getSingleDecl())) {
-            if (isNSAutoreleasePool(VD->getType())) {
-              PoolVarInfo &info = PoolVars[VD];
-              info.Dcl = DclS;
-              ReferenceCollector refColl(info.Refs);
-              refColl.lookFor(VD, S);
-              // Does this statement follow the pattern:  
-              // NSAutoreleasePool * pool = [NSAutoreleasePool  new];
-              if (isPoolCreation(VD->getInit())) {
-                Scopes.push_back(PoolScope());
-                Scopes.back().PoolVar = VD;
-                Scopes.back().CompoundParent = S;
-                Scopes.back().Begin = I;
-              }
-            }
-          }
-        }
-      } else if (BinaryOperator *bop = dyn_cast<BinaryOperator>(child)) {
-        if (DeclRefExpr *dref = dyn_cast<DeclRefExpr>(bop->getLHS())) {
-          if (VarDecl *VD = dyn_cast<VarDecl>(dref->getDecl())) {
-            // Does this statement follow the pattern:  
-            // pool = [NSAutoreleasePool  new];
-            if (isNSAutoreleasePool(VD->getType()) &&
-                isPoolCreation(bop->getRHS())) {
-              Scopes.push_back(PoolScope());
-              Scopes.back().PoolVar = VD;
-              Scopes.back().CompoundParent = S;
-              Scopes.back().Begin = I;
-            }
-          }
-        }
-      }
-
-      if (Scopes.empty())
-        continue;
-
-      if (isPoolDrain(Scopes.back().PoolVar, child)) {
-        PoolScope &scope = Scopes.back();
-        scope.End = I;
-        handlePoolScope(scope, S);
-        Scopes.pop_back();
-      }
-    }
-    return true;
-  }
-
-private:
-  void clearUnavailableDiags(Stmt *S) {
-    if (S)
-      Pass.TA.clearDiagnostic(diag::err_unavailable,
-                              diag::err_unavailable_message,
-                              S->getSourceRange