unique_ptr-ify FrontendAction::takeCurrentASTUnit

llvm-svn: 215319

2 files changed, 12 insertions, 9 deletions

diff --git a/clang/include/clang/Frontend/FrontendAction.h b/clang/include/clang/Frontend/FrontendAction.h
index 9ac9d2828f6..95e1c09c0cc 100644
--- a/clang/include/clang/Frontend/FrontendAction.h
+++ b/clang/include/clang/Frontend/FrontendAction.h
@@ -20,6 +20,7 @@
 
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/LangOptions.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/FrontendOptions.h"
 #include "llvm/ADT/StringRef.h"
 #include <memory>
@@ -29,7 +30,6 @@
 namespace clang {
 class ASTConsumer;
 class ASTMergeAction;
-class ASTUnit;
 class CompilerInstance;
 
 /// Abstract base class for actions which can be performed by the frontend.
@@ -146,10 +146,12 @@ public:
     return *CurrentASTUnit;
   }
 
-  ASTUnit *takeCurrentASTUnit() { return CurrentASTUnit.release(); }
+  std::unique_ptr<ASTUnit> takeCurrentASTUnit() {
+    return std::move(CurrentASTUnit);
+  }
 
   void setCurrentInput(const FrontendInputFile &CurrentInput,
-                       ASTUnit *AST = nullptr);
+                       std::unique_ptr<ASTUnit> AST = nullptr);
 
   /// @}
   /// @name Supported Modes
diff --git a/clang/lib/Frontend/FrontendAction.cpp b/clang/lib/Frontend/FrontendAction.cpp
index cf13d7b9ba8..8295d6ddfa9 100644
--- a/clang/lib/Frontend/FrontendAction.cpp
+++ b/clang/lib/Frontend/FrontendAction.cpp
@@ -129,9 +129,9 @@ FrontendAction::FrontendAction() : Instance(nullptr) {}
 FrontendAction::~FrontendAction() {}
 
 void FrontendAction::setCurrentInput(const FrontendInputFile &CurrentInput,
-                                     ASTUnit *AST) {
+                                     std::unique_ptr<ASTUnit> AST) {
   this->CurrentInput = CurrentInput;
-  CurrentASTUnit.reset(AST);
+  CurrentASTUnit = std::move(AST);
 }
 
 ASTConsumer* FrontendAction::CreateWrappedASTConsumer(CompilerInstance &CI,
@@ -189,13 +189,12 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags(&CI.getDiagnostics());
 
-    ASTUnit *AST = ASTUnit::LoadFromASTFile(InputFile, Diags,
-                                            CI.getFileSystemOpts());
+    std::unique_ptr<ASTUnit> AST(
+        ASTUnit::LoadFromASTFile(InputFile, Diags, CI.getFileSystemOpts()));
+
     if (!AST)
       goto failure;
 
-    setCurrentInput(Input, AST);
-
     // Inform the diagnostic client we are processing a source file.
     CI.getDiagnosticClient().BeginSourceFile(CI.getLangOpts(), nullptr);
     HasBegunSourceFile = true;
@@ -207,6 +206,8 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     CI.setPreprocessor(&AST->getPreprocessor());
     CI.setASTContext(&AST->getASTContext());
 
+    setCurrentInput(Input, std::move(AST));
+
     // Initialize the action.
     if (!BeginSourceFileAction(CI, InputFile))
       goto failure;