From b954e98893505d401acddac42650a81b83d8dc39 Mon Sep 17 00:00:00 2001
From: Daniel Dunbar <daniel@zuster.org>
Date: Thu, 25 Feb 2010 04:37:50 +0000
Subject: [PATCH] Frontend: Add CodeGenAction::takeModule().

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@97111 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/clang/Frontend/CodeGenAction.h | 13 ++++++++++
 lib/Frontend/CodeGenAction.cpp         | 34 ++++++++++++++++++++------
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/include/clang/Frontend/CodeGenAction.h b/include/clang/Frontend/CodeGenAction.h
index f782e1287a..07e02e5c6e 100644
--- a/include/clang/Frontend/CodeGenAction.h
+++ b/include/clang/Frontend/CodeGenAction.h
@@ -8,18 +8,31 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Frontend/FrontendAction.h"
+#include "llvm/ADT/OwningPtr.h"
+
+namespace llvm {
+  class Module;
+}
 
 namespace clang {
 
 class CodeGenAction : public ASTFrontendAction {
 private:
   unsigned Act;
+  llvm::OwningPtr<llvm::Module> TheModule;
 
 protected:
   CodeGenAction(unsigned _Act);
 
   virtual ASTConsumer *CreateASTConsumer(CompilerInstance &CI,
                                          llvm::StringRef InFile);
+
+  virtual void EndSourceFileAction();
+
+public:
+  /// takeModule - Take the generated LLVM module, for use after the action has
+  /// been run. The result may be null on failure.
+  llvm::Module *takeModule();
 };
 
 class EmitAssemblyAction : public CodeGenAction {
diff --git a/lib/Frontend/CodeGenAction.cpp b/lib/Frontend/CodeGenAction.cpp
index 7800d1534e..ff24854339 100644
--- a/lib/Frontend/CodeGenAction.cpp
+++ b/lib/Frontend/CodeGenAction.cpp
@@ -62,7 +62,7 @@ namespace {
 
     llvm::OwningPtr<CodeGenerator> Gen;
 
-    llvm::Module *TheModule;
+    llvm::OwningPtr<llvm::Module> TheModule;
     llvm::TargetData *TheTargetData;
 
     mutable FunctionPassManager *CodeGenPasses;
@@ -97,7 +97,7 @@ namespace {
       LLVMIRGeneration("LLVM IR Generation Time"),
       CodeGenerationTime("Code Generation Time"),
       Gen(CreateLLVMCodeGen(Diags, infile, compopts, C)),
-      TheModule(0), TheTargetData(0),
+      TheTargetData(0),
       CodeGenPasses(0), PerModulePasses(0), PerFunctionPasses(0) {
 
       if (AsmOutStream)
@@ -109,12 +109,13 @@ namespace {
 
     ~BackendConsumer() {
       delete TheTargetData;
-      delete TheModule;
       delete CodeGenPasses;
       delete PerModulePasses;
       delete PerFunctionPasses;
     }
 
+    llvm::Module *takeModule() { return TheModule.take(); }
+
     virtual void Initialize(ASTContext &Ctx) {
       Context = &Ctx;
 
@@ -123,7 +124,7 @@ namespace {
 
       Gen->Initialize(Ctx);
 
-      TheModule = Gen->GetModule();
+      TheModule.reset(Gen->GetModule());
       TheTargetData = new llvm::TargetData(Ctx.Target.getTargetDescription());
 
       if (llvm::TimePassesIsEnabled)
@@ -179,7 +180,7 @@ namespace {
 
 FunctionPassManager *BackendConsumer::getCodeGenPasses() const {
   if (!CodeGenPasses) {
-    CodeGenPasses = new FunctionPassManager(TheModule);
+    CodeGenPasses = new FunctionPassManager(&*TheModule);
     CodeGenPasses->add(new TargetData(*TheTargetData));
   }
 
@@ -197,7 +198,7 @@ PassManager *BackendConsumer::getPerModulePasses() const {
 
 FunctionPassManager *BackendConsumer::getPerFunctionPasses() const {
   if (!PerFunctionPasses) {
-    PerFunctionPasses = new FunctionPassManager(TheModule);
+    PerFunctionPasses = new FunctionPassManager(&*TheModule);
     PerFunctionPasses->add(new TargetData(*TheTargetData));
   }
 
@@ -391,11 +392,12 @@ void BackendConsumer::EmitAssembly() {
   if (!M) {
     // The module has been released by IR gen on failures, do not
     // double free.
-    TheModule = 0;
+    TheModule.take();
     return;
   }
 
-  assert(TheModule == M && "Unexpected module change during IR generation");
+  assert(TheModule.get() == M &&
+         "Unexpected module change during IR generation");
 
   CreatePasses();
   if (!AddEmitPasses())
@@ -433,6 +435,22 @@ void BackendConsumer::EmitAssembly() {
 
 CodeGenAction::CodeGenAction(unsigned _Act) : Act(_Act) {}
 
+void CodeGenAction::EndSourceFileAction() {
+  // If the consumer creation failed, do nothing.
+  if (!getCompilerInstance().hasASTConsumer())
+    return;
+
+  // Steal the module from the consumer.
+  BackendConsumer *Consumer = static_cast<BackendConsumer*>(
+    &getCompilerInstance().getASTConsumer());
+
+  TheModule.reset(Consumer->takeModule());
+}
+
+llvm::Module *CodeGenAction::takeModule() {
+  return TheModule.take();
+}
+
 ASTConsumer *CodeGenAction::CreateASTConsumer(CompilerInstance &CI,
                                               llvm::StringRef InFile) {
   BackendAction BA = static_cast<BackendAction>(Act);
-- 
2.40.0