3 files changed, 49 insertions, 3 deletions
diff --git a/lib/ARCMigrate/ARCMT.cpp b/lib/ARCMigrate/ARCMT.cpp
index cd4835ef39..f9a1a75d87 100644
--- a/lib/ARCMigrate/ARCMT.cpp
+++ b/lib/ARCMigrate/ARCMT.cpp
@@ -232,7 +232,8 @@ bool arcmt::checkForManualIssues(CompilerInvocation &origCI,
   bool NoNSAllocReallocError = origCI.getMigratorOpts().NoNSAllocReallocError;
   bool NoFinalizeRemoval = origCI.getMigratorOpts().NoFinalizeRemoval;
 
-  std::vector<TransformFn> transforms = arcmt::getAllTransformations(OrigGCMode);
+  std::vector<TransformFn> transforms = arcmt::getAllTransformations(OrigGCMode,
+                                                                     NoFinalizeRemoval);
   assert(!transforms.empty());
 
   llvm::OwningPtr<CompilerInvocation> CInvok;
@@ -337,8 +338,10 @@ static bool applyTransforms(CompilerInvocation &origCI,
   CInvok.getFrontendOpts().Inputs.push_back(Input);
   
   MigrationProcess migration(CInvok, DiagClient, outputDir);
+  bool NoFinalizeRemoval = origCI.getMigratorOpts().NoFinalizeRemoval;
 
-  std::vector<TransformFn> transforms = arcmt::getAllTransformations(OrigGCMode);
+  std::vector<TransformFn> transforms = arcmt::getAllTransformations(OrigGCMode,
+                                                                     NoFinalizeRemoval);
   assert(!transforms.empty());
 
   for (unsigned i=0, e = transforms.size(); i != e; ++i) {
diff --git a/lib/ARCMigrate/Transforms.cpp b/lib/ARCMigrate/Transforms.cpp
index a64124015c..91121f0245 100644
--- a/lib/ARCMigrate/Transforms.cpp
+++ b/lib/ARCMigrate/Transforms.cpp
@@ -502,6 +502,45 @@ void MigrationContext::traverse(TranslationUnitDecl *TU) {
   ASTTransform(*this).TraverseDecl(TU);
 }
 
+static void GCRewriteFinalize(MigrationPass &pass) {
+  ASTContext &Ctx = pass.Ctx;
+  TransformActions &TA = pass.TA;
+  DeclContext *DC = Ctx.getTranslationUnitDecl();
+  Selector FinalizeSel =
+   Ctx.Selectors.getNullarySelector(&pass.Ctx.Idents.get("finalize"));
+  
+  typedef DeclContext::specific_decl_iterator<ObjCImplementationDecl>
+  impl_iterator;
+  for (impl_iterator I = impl_iterator(DC->decls_begin()),
+       E = impl_iterator(DC->decls_end()); I != E; ++I) {
+    for (ObjCImplementationDecl::instmeth_iterator
+         MI = (*I)->instmeth_begin(),
+         ME = (*I)->instmeth_end(); MI != ME; ++MI) {
+      ObjCMethodDecl *MD = *MI;
+      if (!MD->hasBody())
+        continue;
+      
+      if (MD->isInstanceMethod() && MD->getSelector() == FinalizeSel) {
+        ObjCMethodDecl *FinalizeM = MD;
+        Transaction Trans(TA);
+        TA.insert(FinalizeM->getSourceRange().getBegin(), 
+                  "#if !__has_feature(objc_arc)\n");
+        CharSourceRange::getTokenRange(FinalizeM->getSourceRange());
+        const SourceManager &SM = pass.Ctx.getSourceManager();
+        const LangOptions &LangOpts = pass.Ctx.getLangOptions();
+        bool Invalid;
+        std::string str = "\n#endif\n";
+        str += Lexer::getSourceText(
+                  CharSourceRange::getTokenRange(FinalizeM->getSourceRange()), 
+                                    SM, LangOpts, &Invalid);
+        TA.insertAfterToken(FinalizeM->getSourceRange().getEnd(), str);
+        
+        break;
+      }
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // getAllTransformations.
 //===----------------------------------------------------------------------===//
@@ -531,9 +570,12 @@ static void independentTransforms(MigrationPass &pass) {
 }
 
 std::vector<TransformFn> arcmt::getAllTransformations(
-                                               LangOptions::GCMode OrigGCMode) {
+                                               LangOptions::GCMode OrigGCMode,
+                                               bool NoFinalizeRemoval) {
   std::vector<TransformFn> transforms;
 
+  if (OrigGCMode ==  LangOptions::GCOnly && NoFinalizeRemoval)
+    transforms.push_back(GCRewriteFinalize);
   transforms.push_back(independentTransforms);
   // This depends on previous transformations removing various expressions.
   transforms.push_back(removeEmptyStatementsAndDeallocFinalize);
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
index 5a0117af5a..a189155e12 100644
--- a/lib/Frontend/CompilerInvocation.cpp
+++ b/lib/Frontend/CompilerInvocation.cpp
@@ -1065,6 +1065,7 @@ static bool ParseAnalyzerArgs(AnalyzerOptions &Opts, ArgList &Args,
 
 static bool ParseMigratorArgs(MigratorOptions &Opts, ArgList &Args) {
   Opts.NoNSAllocReallocError = Args.hasArg(OPT_migrator_no_nsalloc_error);
+  Opts.NoFinalizeRemoval = Args.hasArg(OPT_migrator_no_finalize_removal);
   return true;
 }