3 files changed, 35 insertions, 4 deletions
diff --git a/lib/ARCMigrate/TransProperties.cpp b/lib/ARCMigrate/TransProperties.cpp
index db34ee27a5..6f9166fb98 100644
--- a/lib/ARCMigrate/TransProperties.cpp
+++ b/lib/ARCMigrate/TransProperties.cpp
@@ -526,6 +526,8 @@ public:
 
 } // anonymous namespace
 
-void trans::rewriteProperties(MigrationPass &pass) {
-  ImplementationChecker(pass).TraverseDecl(pass.Ctx.getTranslationUnitDecl());
+void PropertyRewriteTraverser::traverseObjCImplementation(
+                                           ObjCImplementationContext &ImplCtx) {
+  PropertiesRewriter(ImplCtx.getMigrationContext().getPass())
+                                  .doTransform(ImplCtx.getImplementationDecl());
 }
diff --git a/lib/ARCMigrate/Transforms.cpp b/lib/ARCMigrate/Transforms.cpp
index d0c9bc69c8..5ea454a917 100644
--- a/lib/ARCMigrate/Transforms.cpp
+++ b/lib/ARCMigrate/Transforms.cpp
@@ -302,6 +302,16 @@ class ASTTransform : public RecursiveASTVisitor<ASTTransform> {
 public:
   ASTTransform(MigrationContext &MigrateCtx) : MigrateCtx(MigrateCtx) { }
 
+  bool TraverseObjCImplementationDecl(ObjCImplementationDecl *D) {
+    ObjCImplementationContext ImplCtx(MigrateCtx, D);
+    for (MigrationContext::traverser_iterator
+           I = MigrateCtx.traversers_begin(),
+           E = MigrateCtx.traversers_end(); I != E; ++I)
+      (*I)->traverseObjCImplementation(ImplCtx);
+
+    return true;
+  }
+
   bool TraverseStmt(Stmt *rootS) {
     if (!rootS)
       return true;
@@ -358,13 +368,13 @@ static void traverseAST(MigrationPass &pass) {
   if (pass.isGCMigration()) {
     MigrateCtx.addTraverser(new GCCollectableCallsTraverser);
   }
+  MigrateCtx.addTraverser(new PropertyRewriteTraverser());
 
   MigrateCtx.traverse(pass.Ctx.getTranslationUnitDecl());
 }
 
 static void independentTransforms(MigrationPass &pass) {
   rewriteAutoreleasePool(pass);
-  rewriteProperties(pass);
   removeRetainReleaseDeallocFinalize(pass);
   rewriteUnusedInitDelegate(pass);
   removeZeroOutPropsInDeallocFinalize(pass);
diff --git a/lib/ARCMigrate/Transforms.h b/lib/ARCMigrate/Transforms.h
index 515844656a..4ff35a5daf 100644
--- a/lib/ARCMigrate/Transforms.h
+++ b/lib/ARCMigrate/Transforms.h
@@ -37,7 +37,6 @@ void rewriteUnbridgedCasts(MigrationPass &pass);
 void makeAssignARCSafe(MigrationPass &pass);
 void removeRetainReleaseDeallocFinalize(MigrationPass &pass);
 void removeZeroOutPropsInDeallocFinalize(MigrationPass &pass);
-void rewriteProperties(MigrationPass &pass);
 void rewriteBlockObjCVariable(MigrationPass &pass);
 void rewriteUnusedInitDelegate(MigrationPass &pass);
 void checkAPIUses(MigrationPass &pass);
@@ -58,10 +57,25 @@ public:
   Stmt *getTopStmt() { return TopStmt; }
 };
 
+class ObjCImplementationContext {
+  MigrationContext &MigrateCtx;
+  ObjCImplementationDecl *ImpD;
+
+public:
+  ObjCImplementationContext(MigrationContext &MigrateCtx,
+                            ObjCImplementationDecl *D)
+    : MigrateCtx(MigrateCtx), ImpD(D) {}
+
+  MigrationContext &getMigrationContext() { return MigrateCtx; }
+  ObjCImplementationDecl *getImplementationDecl() { return ImpD; }
+};
+
 class ASTTraverser {
 public:
   virtual ~ASTTraverser();
+  virtual void traverseTU(MigrationContext &MigrateCtx) { }
   virtual void traverseBody(BodyContext &BodyCtx) { }
+  virtual void traverseObjCImplementation(ObjCImplementationContext &ImplCtx) {}
 };
 
 class MigrationContext {
@@ -87,6 +101,11 @@ public:
   void traverse(TranslationUnitDecl *TU);
 };
 
+class PropertyRewriteTraverser : public ASTTraverser {
+public:
+  virtual void traverseObjCImplementation(ObjCImplementationContext &ImplCtx);
+};
+
 // GC transformations
 
 class GCCollectableCallsTraverser : public ASTTraverser {