1 files changed, 57 insertions, 10 deletions

diff --git a/lib/ARCMigrate/ARCMT.cpp b/lib/ARCMigrate/ARCMT.cpp
index 9354dc3..f291dec 100644
--- a/lib/ARCMigrate/ARCMT.cpp
+++ b/lib/ARCMigrate/ARCMT.cpp
@@ -91,11 +91,40 @@ namespace {
 
 class CaptureDiagnosticConsumer : public DiagnosticConsumer {
   DiagnosticsEngine &Diags;
+  DiagnosticConsumer &DiagClient;
   CapturedDiagList &CapturedDiags;
+  bool HasBegunSourceFile;
 public:
   CaptureDiagnosticConsumer(DiagnosticsEngine &diags,
-                           CapturedDiagList &capturedDiags)
-    : Diags(diags), CapturedDiags(capturedDiags) { }
+                            DiagnosticConsumer &client,
+                            CapturedDiagList &capturedDiags)
+    : Diags(diags), DiagClient(client), CapturedDiags(capturedDiags),
+      HasBegunSourceFile(false) { }
+
+  virtual void BeginSourceFile(const LangOptions &Opts,
+                               const Preprocessor *PP) {
+    // Pass BeginSourceFile message onto DiagClient on first call.
+    // The corresponding EndSourceFile call will be made from an
+    // explicit call to FinishCapture.
+    if (!HasBegunSourceFile) {
+      DiagClient.BeginSourceFile(Opts, PP);
+      HasBegunSourceFile = true;
+    }
+  }
+
+  void FinishCapture() {
+    // Call EndSourceFile on DiagClient on completion of capture to
+    // enable VerifyDiagnosticConsumer to check diagnostics *after*
+    // it has received the diagnostic list.
+    if (HasBegunSourceFile) {
+      DiagClient.EndSourceFile();
+      HasBegunSourceFile = false;
+    }
+  }
+
+  virtual ~CaptureDiagnosticConsumer() {
+    assert(!HasBegunSourceFile && "FinishCapture not called!");
+  }
 
   virtual void HandleDiagnostic(DiagnosticsEngine::Level level,
                                 const Diagnostic &Info) {
@@ -195,8 +224,19 @@ createInvocationForMigration(CompilerInvocation &origCI) {
   CInvok->getLangOpts()->ObjCAutoRefCount = true;
   CInvok->getLangOpts()->setGC(LangOptions::NonGC);
   CInvok->getDiagnosticOpts().ErrorLimit = 0;
-  CInvok->getDiagnosticOpts().Warnings.push_back(
-                                            "error=arc-unsafe-retained-assign");
+  CInvok->getDiagnosticOpts().PedanticErrors = 0;
+
+  // Ignore -Werror flags when migrating.
+  std::vector<std::string> WarnOpts;
+  for (std::vector<std::string>::iterator
+         I = CInvok->getDiagnosticOpts().Warnings.begin(),
+         E = CInvok->getDiagnosticOpts().Warnings.end(); I != E; ++I) {
+    if (!StringRef(*I).startswith("error"))
+      WarnOpts.push_back(*I);
+  }
+  WarnOpts.push_back("error=arc-unsafe-retained-assign");
+  CInvok->getDiagnosticOpts().Warnings = llvm_move(WarnOpts);
+
   CInvok->getLangOpts()->ObjCRuntimeHasWeak = HasARCRuntime(origCI);
 
   return CInvok.take();
@@ -249,13 +289,15 @@ bool arcmt::checkForManualIssues(CompilerInvocation &origCI,
       new DiagnosticsEngine(DiagID, DiagClient, /*ShouldOwnClient=*/false));
 
   // Filter of all diagnostics.
-  CaptureDiagnosticConsumer errRec(*Diags, capturedDiags);
+  CaptureDiagnosticConsumer errRec(*Diags, *DiagClient, capturedDiags);
   Diags->setClient(&errRec, /*ShouldOwnClient=*/false);
 
   OwningPtr<ASTUnit> Unit(
       ASTUnit::LoadFromCompilerInvocationAction(CInvok.take(), Diags));
-  if (!Unit)
+  if (!Unit) {
+    errRec.FinishCapture();
     return true;
+  }
 
   // Don't filter diagnostics anymore.
   Diags->setClient(DiagClient, /*ShouldOwnClient=*/false);
@@ -267,6 +309,7 @@ bool arcmt::checkForManualIssues(CompilerInvocation &origCI,
     DiagClient->BeginSourceFile(Ctx.getLangOpts(), &Unit->getPreprocessor());
     capturedDiags.reportDiagnostics(*Diags);
     DiagClient->EndSourceFile();
+    errRec.FinishCapture();
     return true;
   }
 
@@ -304,6 +347,7 @@ bool arcmt::checkForManualIssues(CompilerInvocation &origCI,
   capturedDiags.reportDiagnostics(*Diags);
 
   DiagClient->EndSourceFile();
+  errRec.FinishCapture();
 
   // If we are migrating code that gets the '-fobjc-arc' flag, make sure
   // to remove it so that we don't get errors from normal compilation.
@@ -480,13 +524,12 @@ public:
 
 class RewritesApplicator : public TransformActions::RewriteReceiver {
   Rewriter &rewriter;
-  ASTContext &Ctx;
   MigrationProcess::RewriteListener *Listener;
 
 public:
   RewritesApplicator(Rewriter &rewriter, ASTContext &ctx,
                      MigrationProcess::RewriteListener *listener)
-    : rewriter(rewriter), Ctx(ctx), Listener(listener) {
+    : rewriter(rewriter), Listener(listener) {
     if (Listener)
       Listener->start(ctx);
   }
@@ -553,7 +596,7 @@ bool MigrationProcess::applyTransform(TransformFn trans,
       new DiagnosticsEngine(DiagID, DiagClient, /*ShouldOwnClient=*/false));
 
   // Filter of all diagnostics.
-  CaptureDiagnosticConsumer errRec(*Diags, capturedDiags);
+  CaptureDiagnosticConsumer errRec(*Diags, *DiagClient, capturedDiags);
   Diags->setClient(&errRec, /*ShouldOwnClient=*/false);
 
   OwningPtr<ARCMTMacroTrackerAction> ASTAction;
@@ -562,8 +605,10 @@ bool MigrationProcess::applyTransform(TransformFn trans,
   OwningPtr<ASTUnit> Unit(
       ASTUnit::LoadFromCompilerInvocationAction(CInvok.take(), Diags,
                                                 ASTAction.get()));
-  if (!Unit)
+  if (!Unit) {
+    errRec.FinishCapture();
     return true;
+  }
   Unit->setOwnsRemappedFileBuffers(false); // FileRemapper manages that.
 
   // Don't filter diagnostics anymore.
@@ -576,6 +621,7 @@ bool MigrationProcess::applyTransform(TransformFn trans,
     DiagClient->BeginSourceFile(Ctx.getLangOpts(), &Unit->getPreprocessor());
     capturedDiags.reportDiagnostics(*Diags);
     DiagClient->EndSourceFile();
+    errRec.FinishCapture();
     return true;
   }
 
@@ -599,6 +645,7 @@ bool MigrationProcess::applyTransform(TransformFn trans,
   }
 
   DiagClient->EndSourceFile();
+  errRec.FinishCapture();
 
   if (DiagClient->getNumErrors())
     return true;