1 files changed, 43 insertions, 1 deletions
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);