The ARC Migration Tool. All the credit goes to Argyrios and Fariborz

for this.



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

1 files changed, 2094 insertions, 0 deletions

diff --git a/lib/ARCMigrate/Transforms.cpp b/lib/ARCMigrate/Transforms.cpp
new file mode 100644
index 0000000000..ab3bc1dbe7
--- /dev/null
+++ b/lib/ARCMigrate/Transforms.cpp
@@ -0,0 +1,2094 @@
+//===--- Tranforms.cpp - Tranformations to ARC mode -----------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// 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 "Internals.h"
+#include "clang/Sema/SemaDiagnostic.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/AST/StmtVisitor.h"
+#include "clang/AST/ParentMap.h"
+#include "clang/Analysis/DomainSpecific/CocoaConventions.h"
+#include "clang/Lex/Lexer.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/DenseSet.h"
+#include <map>
+
+using namespace clang;
+using namespace arcmt;
+using llvm::StringRef;
+
+//===----------------------------------------------------------------------===//
+// Transformations.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class RemovablesCollector : public RecursiveASTVisitor<RemovablesCollector> {
+  llvm::DenseSet<Expr *> &Removables;
+
+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);
+  }
+};
+
+} // end anonymous namespace.
+
+static bool HasSideEffects(Expr *E, ASTContext &Ctx) {
+  if (!E || !E->HasSideEffects(Ctx))
+    return false;
+
+  E = E->IgnoreParenCasts();
+  ObjCMessageExpr *ME = dyn_cast<ObjCMessageExpr>(E);
+  if (!ME)
+    return true;
+  switch (ME->getMethodFamily()) {
+  case OMF_autorelease:
+  case OMF_dealloc:
+  case OMF_release:
+  case OMF_retain:
+    switch (ME->getReceiverKind()) {
+    case ObjCMessageExpr::SuperInstance:
+      return false;
+    case ObjCMessageExpr::Instance:
+      return HasSideEffects(ME->getInstanceReceiver(), Ctx);
+    default:
+      break;
+    }
+    break;
+  default:
+    break;
+  }
+
+  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;
+public:
+  ReferenceClear(llvm::DenseSet<Expr *> &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;
+
+public:
+  ReferenceCollector(llvm::DenseSet<Expr *> &refs)
+    : Dcl(0), Refs(refs) { }
+
+  void lookFor(ValueDecl *D, Stmt *S) {
+    Dcl = D;
+    TraverseStmt(S);
+  }
+
+  bool VisitDeclRefExpr(DeclRefExpr *E) {
+    if (E->getDecl() == Dcl)
+      Refs.insert(E);
+    return true;
+  }
+
+  bool VisitBlockDeclRefExpr(BlockDeclRefExpr *E) {
+    if (E->getDecl() == Dcl)
+      Refs.insert(E);
+    return true;
+  }
+};
+
+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->getType().getLocalQualifiers().getObjCLifetime()
+          == Qualifiers::OCL_ExplicitNone &&
+          (var->getTypeSourceInfo() &&
+           !var->getTypeSourceInfo()->getType().isConstQualified())) {
+        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;
+public:
+  NonObjCToObjCCaster(MigrationPass &pass) : Pass(pass) { }
+
+  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) {
+    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;
+
+    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(), retained ? "__bridge_transfer "
+                                                        : "__bridge ");
+    } else {
+      SourceLocation insertLoc = E->getSubExpr()->getLocStart();
+      llvm::SmallString<128> newCast;
+      newCast += '(';
+      newCast +=  retained ? "__bridge_transfer " : "__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) {
+    // FIXME: Handle these casts.
+    return;
+#if 0
+    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;
+
+    Transaction Trans(TA);
+    TA.clearDiagnostic(diag::err_arc_mismatched_cast,
+                              diag::err_arc_cast_requires_bridge,
+                              E->getLocStart());
+
+    assert(!E->getType()->isObjCObjectPointerType());
+
+    bool shouldCast = !isa<CStyleCastExpr>(E) &&
+                      !E->getType()->getPointeeType().isConstQualified();
+    SourceLocation loc = E->getSubExpr()->getLocStart();
+    if (isa<ParenExpr>(E->getSubExpr())) {
+      TA.insert(loc, shouldCast ? "(void*)objc_unretainedPointer"
+                                : "objc_unretainedPointer");
+    } else {
+      TA.insert(loc, shouldCast ? "(void*)objc_unretainedPointer("
+                                : "objc_unretainedPointer(");
+      TA.insertAfterToken(E->getLocEnd(), ")");
+    }
+#endif
+  }
+
+  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 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());
+  }
+
+  struct PoolScope {
+    VarDecl *PoolVar;
+    CompoundStmt *CompoundParent;
+    Stmt::child_iterator Begin;
+    Stmt::child_iterator End;
+    bool IsFollowedBySimpleReturnStmt;
+    llvm::SmallVector<ObjCMessageExpr *, 4> Releases;
+
+    PoolScope() : PoolVar(0), CompoundParent(0), Begin(), End(),