1 files changed, 4 insertions, 1 deletions
diff --git a/lib/Rewrite/FixItRewriter.cpp b/lib/Rewrite/FixItRewriter.cpp
index 3fdec07459..582096978d 100644
--- a/lib/Rewrite/FixItRewriter.cpp
+++ b/lib/Rewrite/FixItRewriter.cpp
@@ -32,11 +32,12 @@ FixItRewriter::FixItRewriter(Diagnostic &Diags, SourceManager &SourceMgr,
     Rewrite(SourceMgr, LangOpts),
     FixItOpts(FixItOpts),
     NumFailures(0) {
-  Client = Diags.getClient();
+  Client = Diags.takeClient();
   Diags.setClient(this);
 }
 
 FixItRewriter::~FixItRewriter() {
+  Diags.takeClient();
   Diags.setClient(Client);
 }
 
@@ -144,9 +145,11 @@ void FixItRewriter::Diag(FullSourceLoc Loc, unsigned DiagID) {
   // When producing this diagnostic, we temporarily bypass ourselves,
   // clear out any current diagnostic, and let the downstream client
   // format the diagnostic.
+  Diags.takeClient();
   Diags.setClient(Client);
   Diags.Clear();
   Diags.Report(Loc, DiagID);
+  Diags.takeClient();
   Diags.setClient(this);
 }