1 files changed, 6 insertions, 9 deletions
diff --git a/clang/lib/Frontend/Rewrite/FixItRewriter.cpp b/clang/lib/Frontend/Rewrite/FixItRewriter.cpp
index f1776239d53..a3e14f9f020 100644
--- a/clang/lib/Frontend/Rewrite/FixItRewriter.cpp
+++ b/clang/lib/Frontend/Rewrite/FixItRewriter.cpp
@@ -36,14 +36,13 @@ FixItRewriter::FixItRewriter(DiagnosticsEngine &Diags, SourceManager &SourceMgr,
     FixItOpts(FixItOpts),
     NumFailures(0),
     PrevDiagSilenced(false) {
-  OwnsClient = Diags.ownsClient();
-  Client = Diags.takeClient();
-  Diags.setClient(this);
+  Owner = Diags.takeClient();
+  Client = Diags.getClient();
+  Diags.setClient(this, false);
 }
 
 FixItRewriter::~FixItRewriter() {
-  Diags.takeClient();
-  Diags.setClient(Client, OwnsClient);
+  Diags.setClient(Client, Owner.release() != nullptr);
 }
 
 bool FixItRewriter::WriteFixedFile(FileID ID, raw_ostream &OS) {
@@ -188,12 +187,10 @@ void FixItRewriter::Diag(SourceLocation 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.setClient(Client, false);
   Diags.Clear();
   Diags.Report(Loc, DiagID);
-  Diags.takeClient();
-  Diags.setClient(this);
+  Diags.setClient(this, false);
 }
 
 FixItOptions::~FixItOptions() {}