libLTO: Allow LTOModule to own a context

llvm-svn: 221728

3 files changed, 59 insertions, 9 deletions

diff --git a/llvm/include/llvm/LTO/LTOCodeGenerator.h b/llvm/include/llvm/LTO/LTOCodeGenerator.h
index f93f013b7c6..0c9ce4a54f0 100644
--- a/llvm/include/llvm/LTO/LTOCodeGenerator.h
+++ b/llvm/include/llvm/LTO/LTOCodeGenerator.h
@@ -119,6 +119,8 @@ struct LTOCodeGenerator {
 
   void setDiagnosticHandler(lto_diagnostic_handler_t, void *);
 
+  LLVMContext &getContext() { return Context; }
+
 private:
   void initializeLTOPasses();
 
diff --git a/llvm/include/llvm/LTO/LTOModule.h b/llvm/include/llvm/LTO/LTOModule.h
index fc583a20f68..23b1fce435b 100644
--- a/llvm/include/llvm/LTO/LTOModule.h
+++ b/llvm/include/llvm/LTO/LTOModule.h
@@ -46,6 +46,8 @@ private:
     const GlobalValue *symbol;
   };
 
+  std::unique_ptr<LLVMContext> OwnedContext;
+
   std::unique_ptr<object::IRObjectFile> IRFile;
   std::unique_ptr<TargetMachine> _target;
   StringSet                               _linkeropt_strings;
@@ -59,8 +61,12 @@ private:
   std::vector<const char*>                _asm_undefines;
 
   LTOModule(std::unique_ptr<object::IRObjectFile> Obj, TargetMachine *TM);
+  LTOModule(std::unique_ptr<object::IRObjectFile> Obj, TargetMachine *TM,
+            std::unique_ptr<LLVMContext> Context);
 
 public:
+  ~LTOModule();
+
   /// Returns 'true' if the file or memory contents is LLVM bitcode.
   static bool isBitcodeFile(const void *mem, size_t length);
   static bool isBitcodeFile(const char *path);
@@ -95,6 +101,13 @@ public:
                                      TargetOptions options, std::string &errMsg,
                                      StringRef path = "");
 
+  static LTOModule *createInLocalContext(const void *mem, size_t length,
+                                         TargetOptions options,
+                                         std::string &errMsg, StringRef path);
+  static LTOModule *createInContext(const void *mem, size_t length,
+                                    TargetOptions options, std::string &errMsg,
+                                    StringRef path, LLVMContext *Context);
+
   const Module &getModule() const {
     return const_cast<LTOModule*>(this)->getModule();
   }
@@ -204,7 +217,7 @@ private:
 
   /// Create an LTOModule (private version).
   static LTOModule *makeLTOModule(MemoryBufferRef Buffer, TargetOptions options,
-                                  std::string &errMsg);
+                                  std::string &errMsg, LLVMContext *Context);
 };
 }
 #endif
diff --git a/llvm/lib/LTO/LTOModule.cpp b/llvm/lib/LTO/LTOModule.cpp
index b9320d24cc3..34e7f51c7bf 100644
--- a/llvm/lib/LTO/LTOModule.cpp
+++ b/llvm/lib/LTO/LTOModule.cpp
@@ -52,6 +52,13 @@ LTOModule::LTOModule(std::unique_ptr<object::IRObjectFile> Obj,
                      llvm::TargetMachine *TM)
     : IRFile(std::move(Obj)), _target(TM) {}
 
+LTOModule::LTOModule(std::unique_ptr<object::IRObjectFile> Obj,
+                     llvm::TargetMachine *TM,
+                     std::unique_ptr<LLVMContext> Context)
+    : OwnedContext(std::move(Context)), IRFile(std::move(Obj)), _target(TM) {}
+
+LTOModule::~LTOModule() {}
+
 /// isBitcodeFile - Returns 'true' if the file (or memory contents) is LLVM
 /// bitcode.
 bool LTOModule::isBitcodeFile(const void *Mem, size_t Length) {
@@ -77,7 +84,8 @@ bool LTOModule::isBitcodeForTarget(MemoryBuffer *Buffer,
       IRObjectFile::findBitcodeInMemBuffer(Buffer->getMemBufferRef());
   if (!BCOrErr)
     return false;
-  std::string Triple = getBitcodeTargetTriple(*BCOrErr, getGlobalContext());
+  LLVMContext Context;
+  std::string Triple = getBitcodeTargetTriple(*BCOrErr, Context);
   return StringRef(Triple).startswith(TriplePrefix);
 }
 
@@ -90,7 +98,8 @@ LTOModule *LTOModule::createFromFile(const char *path, TargetOptions options,
     return nullptr;
   }
   std::unique_ptr<MemoryBuffer> Buffer = std::move(BufferOrErr.get());
-  return makeLTOModule(Buffer->getMemBufferRef(), options, errMsg);
+  return makeLTOModule(Buffer->getMemBufferRef(), options, errMsg,
+                       &getGlobalContext());
 }
 
 LTOModule *LTOModule::createFromOpenFile(int fd, const char *path, size_t size,
@@ -110,27 +119,49 @@ LTOModule *LTOModule::createFromOpenFileSlice(int fd, const char *path,
     return nullptr;
   }
   std::unique_ptr<MemoryBuffer> Buffer = std::move(BufferOrErr.get());
-  return makeLTOModule(Buffer->getMemBufferRef(), options, errMsg);
+  return makeLTOModule(Buffer->getMemBufferRef(), options, errMsg,
+                       &getGlobalContext());
 }
 
 LTOModule *LTOModule::createFromBuffer(const void *mem, size_t length,
                                        TargetOptions options,
                                        std::string &errMsg, StringRef path) {
+  return createInContext(mem, length, options, errMsg, path,
+                         &getGlobalContext());
+}
+
+LTOModule *LTOModule::createInLocalContext(const void *mem, size_t length,
+                                           TargetOptions options,
+                                           std::string &errMsg,
+                                           StringRef path) {
+  return createInContext(mem, length, options, errMsg, path, nullptr);
+}
+
+LTOModule *LTOModule::createInContext(const void *mem, size_t length,
+                                      TargetOptions options,
+                                      std::string &errMsg, StringRef path,
+                                      LLVMContext *Context) {
   StringRef Data((const char *)mem, length);
   MemoryBufferRef Buffer(Data, path);
-  return makeLTOModule(Buffer, options, errMsg);
+  return makeLTOModule(Buffer, options, errMsg, Context);
 }
 
 LTOModule *LTOModule::makeLTOModule(MemoryBufferRef Buffer,
-                                    TargetOptions options,
-                                    std::string &errMsg) {
+                                    TargetOptions options, std::string &errMsg,
+                                    LLVMContext *Context) {
+  std::unique_ptr<LLVMContext> OwnedContext;
+  if (!Context) {
+    OwnedContext = llvm::make_unique<LLVMContext>();
+    Context = OwnedContext.get();
+  }
+
   ErrorOr<MemoryBufferRef> MBOrErr =
       IRObjectFile::findBitcodeInMemBuffer(Buffer);
   if (std::error_code EC = MBOrErr.getError()) {
     errMsg = EC.message();
     return nullptr;
   }
-  ErrorOr<Module *> MOrErr = parseBitcodeFile(*MBOrErr, getGlobalContext());
+  ErrorOr<Module *> MOrErr = parseBitcodeFile(*MBOrErr, *Context);
   if (std::error_code EC = MOrErr.getError()) {
     errMsg = EC.message();
     return nullptr;
@@ -169,7 +200,11 @@ LTOModule *LTOModule::makeLTOModule(MemoryBufferRef Buffer,
   std::unique_ptr<object::IRObjectFile> IRObj(
       new object::IRObjectFile(Buffer, std::move(M)));
 
-  LTOModule *Ret = new LTOModule(std::move(IRObj), target);
+  LTOModule *Ret;
+  if (OwnedContext)
+    Ret = new LTOModule(std::move(IRObj), target, std::move(OwnedContext));
+  else
+    Ret = new LTOModule(std::move(IRObj), target);
 
   if (Ret->parseSymbols(errMsg)) {
     delete Ret;