Change linkInModule to take a std::unique_ptr.

Passing in a std::unique_ptr should help find errors when the module
is used after being linked into another module.

llvm-svn: 255842

3 files changed, 44 insertions, 21 deletions

diff --git a/llvm/lib/LTO/LTOCodeGenerator.cpp b/llvm/lib/LTO/LTOCodeGenerator.cpp
index 525ca37c2f1..e51366c2930 100644
--- a/llvm/lib/LTO/LTOCodeGenerator.cpp
+++ b/llvm/lib/LTO/LTOCodeGenerator.cpp
@@ -106,7 +106,7 @@ bool LTOCodeGenerator::addModule(LTOModule *Mod) {
   assert(&Mod->getModule().getContext() == &Context &&
          "Expected module in same context");
 
-  bool ret = IRLinker->linkInModule(Mod->getModule());
+  bool ret = IRLinker->linkInModule(Mod->takeModule());
 
   const std::vector<const char *> &undefs = Mod->getAsmUndefinedRefs();
   for (int i = 0, e = undefs.size(); i != e; ++i)
diff --git a/llvm/lib/Linker/LinkModules.cpp b/llvm/lib/Linker/LinkModules.cpp
index 44b93696be1..1a10796a605 100644
--- a/llvm/lib/Linker/LinkModules.cpp
+++ b/llvm/lib/Linker/LinkModules.cpp
@@ -789,10 +789,15 @@ bool ModuleLinker::run() {
 
 Linker::Linker(Module &M) : Mover(M) {}
 
-bool Linker::linkInModule(Module &Src, unsigned Flags,
+bool Linker::linkInModule(std::unique_ptr<Module> Src, unsigned Flags,
                           const FunctionInfoIndex *Index,
                           DenseSet<const GlobalValue *> *FunctionsToImport) {
-  ModuleLinker TheLinker(Mover, Src, Flags, Index, FunctionsToImport);
+  ModuleLinker TheLinker(Mover, *Src, Flags, Index, FunctionsToImport);
+  return TheLinker.run();
+}
+
+bool Linker::linkInModuleForCAPI(Module &Src) {
+  ModuleLinker TheLinker(Mover, Src, 0, nullptr, nullptr);
   return TheLinker.run();
 }
 
@@ -805,18 +810,19 @@ bool Linker::linkInModule(Module &Src, unsigned Flags,
 /// true is returned and ErrorMsg (if not null) is set to indicate the problem.
 /// Upon failure, the Dest module could be in a modified state, and shouldn't be
 /// relied on to be consistent.
-bool Linker::linkModules(Module &Dest, Module &Src, unsigned Flags) {
+bool Linker::linkModules(Module &Dest, std::unique_ptr<Module> Src,
+                         unsigned Flags) {
   Linker L(Dest);
-  return L.linkInModule(Src, Flags);
+  return L.linkInModule(std::move(Src), Flags);
 }
 
 std::unique_ptr<Module>
-llvm::renameModuleForThinLTO(std::unique_ptr<Module> &M,
+llvm::renameModuleForThinLTO(std::unique_ptr<Module> M,
                              const FunctionInfoIndex *Index) {
   std::unique_ptr<llvm::Module> RenamedModule(
       new llvm::Module(M->getModuleIdentifier(), M->getContext()));
   Linker L(*RenamedModule.get());
-  if (L.linkInModule(*M.get(), llvm::Linker::Flags::None, Index))
+  if (L.linkInModule(std::move(M), llvm::Linker::Flags::None, Index))
     return nullptr;
   return RenamedModule;
 }
@@ -843,7 +849,9 @@ LLVMBool LLVMLinkModules(LLVMModuleRef Dest, LLVMModuleRef Src,
   std::string Message;
   Ctx.setDiagnosticHandler(diagnosticHandler, &Message, true);
 
-  LLVMBool Result = Linker::linkModules(*D, *unwrap(Src));
+  Linker L(*D);
+  Module *M = unwrap(Src);
+  LLVMBool Result = L.linkInModuleForCAPI(*M);
 
   Ctx.setDiagnosticHandler(OldDiagnosticHandler, OldDiagnosticContext, true);
 
@@ -851,3 +859,9 @@ LLVMBool LLVMLinkModules(LLVMModuleRef Dest, LLVMModuleRef Src,
     *OutMessages = strdup(Message.c_str());
   return Result;
 }
+
+LLVMBool LLVMLinkModules2(LLVMModuleRef Dest, LLVMModuleRef Src) {
+  Module *D = unwrap(Dest);
+  std::unique_ptr<Module> M(unwrap(Src));
+  return Linker::linkModules(*D, std::move(M));
+}
diff --git a/llvm/lib/Transforms/IPO/FunctionImport.cpp b/llvm/lib/Transforms/IPO/FunctionImport.cpp
index fe58bdbed19..b8c2ea7dce2 100644
--- a/llvm/lib/Transforms/IPO/FunctionImport.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionImport.cpp
@@ -71,6 +71,14 @@ public:
 
   /// Retrieve a Module from the cache or lazily load it on demand.
   Module &operator()(StringRef FileName);
+
+  std::unique_ptr<Module> takeModule(StringRef FileName) {
+    auto I = ModuleMap.find(FileName);
+    assert(I != ModuleMap.end());
+    std::unique_ptr<Module> Ret = std::move(I->second);
+    ModuleMap.erase(I);
+    return Ret;
+  }
 };
 
 // Get a Module for \p FileName from the cache, or load it lazily.
@@ -149,12 +157,13 @@ static void findExternalCalls(const Module &DestModule, Function &F,
 //
 // \p ModuleToFunctionsToImportMap is filled with the set of Function to import
 // per Module.
-static void GetImportList(
-    Module &DestModule, SmallVector<StringRef, 64> &Worklist,
-    StringSet<> &CalledFunctions,
-    std::map<StringRef, std::pair<Module *, DenseSet<const GlobalValue *>>> &
-        ModuleToFunctionsToImportMap,
-    const FunctionInfoIndex &Index, ModuleLazyLoaderCache &ModuleLoaderCache) {
+static void GetImportList(Module &DestModule,
+                          SmallVector<StringRef, 64> &Worklist,
+                          StringSet<> &CalledFunctions,
+                          std::map<StringRef, DenseSet<const GlobalValue *>>
+                              &ModuleToFunctionsToImportMap,
+                          const FunctionInfoIndex &Index,
+                          ModuleLazyLoaderCache &ModuleLoaderCache) {
   while (!Worklist.empty()) {
     auto CalledFunctionName = Worklist.pop_back_val();
     DEBUG(dbgs() << DestModule.getModuleIdentifier() << ": Process import for "
@@ -238,8 +247,7 @@ static void GetImportList(
 
     // Add the function to the import list
     auto &Entry = ModuleToFunctionsToImportMap[SrcModule.getModuleIdentifier()];
-    Entry.first = &SrcModule;
-    Entry.second.insert(F);
+    Entry.insert(F);
 
     // Process the newly imported functions and add callees to the worklist.
     F->materialize();
@@ -274,7 +282,7 @@ bool FunctionImporter::importFunctions(Module &DestModule) {
   Linker TheLinker(DestModule);
 
   // Map of Module -> List of Function to import from the Module
-  std::map<StringRef, std::pair<Module *, DenseSet<const GlobalValue *>>>
+  std::map<StringRef, DenseSet<const GlobalValue *>>
       ModuleToFunctionsToImportMap;
 
   // Analyze the summaries and get the list of functions to import by
@@ -287,14 +295,15 @@ bool FunctionImporter::importFunctions(Module &DestModule) {
   // Do the actual import of functions now, one Module at a time
   for (auto &FunctionsToImportPerModule : ModuleToFunctionsToImportMap) {
     // Get the module for the import
-    auto &FunctionsToImport = FunctionsToImportPerModule.second.second;
-    auto *SrcModule = FunctionsToImportPerModule.second.first;
+    auto &FunctionsToImport = FunctionsToImportPerModule.second;
+    std::unique_ptr<Module> SrcModule =
+        ModuleLoaderCache.takeModule(FunctionsToImportPerModule.first);
     assert(&DestModule.getContext() == &SrcModule->getContext() &&
            "Context mismatch");
 
     // Link in the specified functions.
-    if (TheLinker.linkInModule(*SrcModule, Linker::Flags::None, &Index,
-                               &FunctionsToImport))
+    if (TheLinker.linkInModule(std::move(SrcModule), Linker::Flags::None,
+                               &Index, &FunctionsToImport))
       report_fatal_error("Function Import: link error");
 
     ImportedCount += FunctionsToImport.size();