3 files changed, 49 insertions, 4 deletions
diff --git a/clang/lib/ARCMigrate/ObjCMT.cpp b/clang/lib/ARCMigrate/ObjCMT.cpp
index 23d61cf27f4..ff2ae40758f 100644
--- a/clang/lib/ARCMigrate/ObjCMT.cpp
+++ b/clang/lib/ARCMigrate/ObjCMT.cpp
@@ -28,6 +28,7 @@
 #include "clang/StaticAnalyzer/Checkers/ObjCRetainCount.h"
 #include "clang/AST/Attr.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/Path.h"
 
 using namespace clang;
 using namespace arcmt;
@@ -90,6 +91,7 @@ public:
   bool IsOutputFile;
   llvm::SmallPtrSet<ObjCProtocolDecl *, 32> ObjCProtocolDecls;
   llvm::SmallVector<const Decl *, 8> CFFunctionIBCandidates;
+  llvm::StringMap<char> WhiteListFilenames;
   
   ObjCMigrateASTConsumer(StringRef migrateDir,
                          unsigned astMigrateActions,
@@ -97,12 +99,19 @@ public:
                          FileManager &fileMgr,
                          const PPConditionalDirectiveRecord *PPRec,
                          Preprocessor &PP,
-                         bool isOutputFile = false)
+                         bool isOutputFile,
+                         ArrayRef<std::string> WhiteList)
   : MigrateDir(migrateDir),
     ASTMigrateActions(astMigrateActions),
     NSIntegerTypedefed(0), NSUIntegerTypedefed(0),
     Remapper(remapper), FileMgr(fileMgr), PPRec(PPRec), PP(PP),
-    IsOutputFile(isOutputFile) { }
+    IsOutputFile(isOutputFile) {
+
+    for (ArrayRef<std::string>::iterator
+           I = WhiteList.begin(), E = WhiteList.end(); I != E; ++I) {
+      WhiteListFilenames.GetOrCreateValue(*I);
+    }
+  }
 
 protected:
   virtual void Initialize(ASTContext &Context) {
@@ -125,6 +134,13 @@ protected:
   }
 
   virtual void HandleTranslationUnit(ASTContext &Ctx);
+
+  bool canModifyFile(StringRef Path) {
+    if (WhiteListFilenames.empty())
+      return true;
+    return WhiteListFilenames.find(llvm::sys::path::filename(Path))
+        != WhiteListFilenames.end();
+  }
 };
 
 }
@@ -151,7 +167,9 @@ ASTConsumer *ObjCMigrateAction::CreateASTConsumer(CompilerInstance &CI,
                                                        Remapper,
                                                     CompInst->getFileManager(),
                                                        PPRec,
-                                                       CompInst->getPreprocessor());
+                                                       CompInst->getPreprocessor(),
+                                                       false,
+                                                       ArrayRef<std::string>());
   ASTConsumer *Consumers[] = { MTConsumer, WrappedConsumer };
   return new MultiplexConsumer(Consumers);
 }
@@ -1682,6 +1700,8 @@ void ObjCMigrateASTConsumer::HandleTranslationUnit(ASTContext &Ctx) {
     assert(file);
     if (IsReallyASystemHeader(Ctx, file, FID))
       continue;
+    if (!canModifyFile(file->getName()))
+      continue;
     SmallString<512> newText;
     llvm::raw_svector_ostream vecOS(newText);
     buf.write(vecOS);
@@ -1705,6 +1725,25 @@ bool MigrateSourceAction::BeginInvocation(CompilerInstance &CI) {
   return true;
 }
 
+static std::vector<std::string> getWhiteListFilenames(StringRef DirPath) {
+  using namespace llvm::sys::fs;
+  using namespace llvm::sys::path;
+
+  std::vector<std::string> Filenames;
+  if (DirPath.empty() || !is_directory(DirPath))
+    return Filenames;
+  
+  llvm::error_code EC;
+  directory_iterator DI = directory_iterator(DirPath, EC);
+  directory_iterator DE;
+  for (; !EC && DI != DE; DI = DI.increment(EC)) {
+    if (is_regular_file(DI->path()))
+      Filenames.push_back(filename(DI->path()));
+  }
+
+  return Filenames;
+}
+
 ASTConsumer *MigrateSourceAction::CreateASTConsumer(CompilerInstance &CI,
                                                   StringRef InFile) {
   PPConditionalDirectiveRecord *
@@ -1721,11 +1760,14 @@ ASTConsumer *MigrateSourceAction::CreateASTConsumer(CompilerInstance &CI,
                     FrontendOptions::ObjCMT_Subscripting;
   }
   CI.getPreprocessor().addPPCallbacks(PPRec);
+  std::vector<std::string> WhiteList =
+    getWhiteListFilenames(CI.getFrontendOpts().ObjCMTWhiteListPath);
   return new ObjCMigrateASTConsumer(CI.getFrontendOpts().OutputFile,
                                     ObjCMTAction,
                                     Remapper,
                                     CI.getFileManager(),
                                     PPRec,
                                     CI.getPreprocessor(),
-                                    /*isOutputFile=*/true); 
+                                    /*isOutputFile=*/true,
+                                    WhiteList);
 }
diff --git a/clang/lib/Driver/Tools.cpp b/clang/lib/Driver/Tools.cpp
index c8ecdc996af..9027f70c3cb 100644
--- a/clang/lib/Driver/Tools.cpp
+++ b/clang/lib/Driver/Tools.cpp
@@ -2726,6 +2726,7 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
     Args.AddLastArg(CmdArgs, options::OPT_objcmt_atomic_property);
     Args.AddLastArg(CmdArgs, options::OPT_objcmt_returns_innerpointer_property);
     Args.AddLastArg(CmdArgs, options::OPT_objcmt_ns_nonatomic_iosonly);
+    Args.AddLastArg(CmdArgs, options::OPT_objcmt_white_list_dir_path);
   }
 
   // Add preprocessing options like -I, -D, etc. if we are using the
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 009d6ec094b..581d56d4621 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -821,6 +821,8 @@ static InputKind ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_objcmt_migrate_all))
     Opts.ObjCMTAction |= FrontendOptions::ObjCMT_MigrateDecls;
 
+  Opts.ObjCMTWhiteListPath = Args.getLastArgValue(OPT_objcmt_white_list_dir_path);
+
   if (Opts.ARCMTAction != FrontendOptions::ARCMT_None &&
       Opts.ObjCMTAction != FrontendOptions::ObjCMT_None) {
     Diags.Report(diag::err_drv_argument_not_allowed_with)