9 files changed, 94 insertions, 5 deletions
diff --git a/llvm/include/llvm-c/lto.h b/llvm/include/llvm-c/lto.h
index 1fe0cd5c75d..c6acdaddca5 100644
--- a/llvm/include/llvm-c/lto.h
+++ b/llvm/include/llvm-c/lto.h
@@ -40,7 +40,7 @@ typedef bool lto_bool_t;
  * @{
  */
 
-#define LTO_API_VERSION 12
+#define LTO_API_VERSION 13
 
 /**
  * \since prior to LTO_API_VERSION=3
@@ -396,6 +396,17 @@ extern lto_bool_t
 lto_codegen_add_module(lto_code_gen_t cg, lto_module_t mod);
 
 /**
+ * Sets the object module for code generation. This will transfer the ownship of
+ * the module to code generator.
+ *
+ * \c cg and \c mod must both be in the same context.
+ *
+ * \since prior to LTO_API_VERSION=13
+ */
+extern void
+lto_codegen_set_module(lto_code_gen_t cg, lto_module_t mod);
+
+/**
  * Sets if debug info should be generated.
  * Returns true on error (check lto_get_error_message() for details).
  *
diff --git a/llvm/include/llvm/LTO/LTOCodeGenerator.h b/llvm/include/llvm/LTO/LTOCodeGenerator.h
index ab307bf2d0c..cdfc543ec24 100644
--- a/llvm/include/llvm/LTO/LTOCodeGenerator.h
+++ b/llvm/include/llvm/LTO/LTOCodeGenerator.h
@@ -67,6 +67,9 @@ struct LTOCodeGenerator {
   // Merge given module, return true on success.
   bool addModule(struct LTOModule *);
 
+  // Set the destination module.
+  void setModule(struct LTOModule *);
+
   void setTargetOptions(TargetOptions options);
   void setDebugInfo(lto_debug_model);
   void setCodePICModel(lto_codegen_model);
diff --git a/llvm/include/llvm/Linker/Linker.h b/llvm/include/llvm/Linker/Linker.h
index 9c3ecea590b..aac9dcdcb36 100644
--- a/llvm/include/llvm/Linker/Linker.h
+++ b/llvm/include/llvm/Linker/Linker.h
@@ -70,6 +70,9 @@ public:
   /// Returns true on error.
   bool linkInModule(Module *Src);
 
+  /// \brief Set the composite to the passed-in module.
+  void setModule(Module *Dst);
+
   static bool LinkModules(Module *Dest, Module *Src,
                           DiagnosticHandlerFunction DiagnosticHandler);
 
diff --git a/llvm/lib/LTO/LTOCodeGenerator.cpp b/llvm/lib/LTO/LTOCodeGenerator.cpp
index 6e5bdaa97c7..1ba4159c9bf 100644
--- a/llvm/lib/LTO/LTOCodeGenerator.cpp
+++ b/llvm/lib/LTO/LTOCodeGenerator.cpp
@@ -141,6 +141,22 @@ bool LTOCodeGenerator::addModule(LTOModule *mod) {
   return !ret;
 }
 
+void LTOCodeGenerator::setModule(LTOModule *Mod) {
+  assert(&Mod->getModule().getContext() == &Context &&
+         "Expected module in same context");
+
+  // Delete the old merged module.
+  if (IRLinker.getModule())
+    IRLinker.deleteModule();
+  AsmUndefinedRefs.clear();
+
+  IRLinker.setModule(&Mod->getModule());
+
+  const std::vector<const char*> &Undefs = Mod->getAsmUndefinedRefs();
+  for (int I = 0, E = Undefs.size(); I != E; ++I)
+    AsmUndefinedRefs[Undefs[I]] = 1;
+}
+
 void LTOCodeGenerator::setTargetOptions(TargetOptions options) {
   Options = options;
 }
diff --git a/llvm/lib/Linker/LinkModules.cpp b/llvm/lib/Linker/LinkModules.cpp
index b55fe483fb2..e6d9acc5098 100644
--- a/llvm/lib/Linker/LinkModules.cpp
+++ b/llvm/lib/Linker/LinkModules.cpp
@@ -1748,6 +1748,10 @@ bool Linker::linkInModule(Module *Src) {
   return RetCode;
 }
 
+void Linker::setModule(Module *Dst) {
+  init(Dst, DiagnosticHandler);
+}
+
 //===----------------------------------------------------------------------===//
 // LinkModules entrypoint.
 //===----------------------------------------------------------------------===//
diff --git a/llvm/test/LTO/X86/set-merged.ll b/llvm/test/LTO/X86/set-merged.ll
new file mode 100644
index 00000000000..0e2e1ea48cf
--- /dev/null
+++ b/llvm/test/LTO/X86/set-merged.ll
@@ -0,0 +1,36 @@
+; RUN: llvm-as < %s >%t1
+; RUN: llvm-lto -exported-symbol=_main -set-merged-module -o %t2 %t1
+; RUN: llvm-objdump -d %t2 | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.10.0"
+
+; CHECK: _main
+; CHECK: movl $132
+define i32 @_Z3fooi(i32 %a) {
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %0 = load i32* %a.addr, align 4
+  %1 = load i32* %a.addr, align 4
+  %call = call i32 @_Z4bar2i(i32 %1)
+  %add = add nsw i32 %0, %call
+  ret i32 %add
+}
+
+define i32 @_Z4bar2i(i32 %a) {
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %0 = load i32* %a.addr, align 4
+  %mul = mul nsw i32 2, %0
+  ret i32 %mul
+}
+
+define i32 @main() {
+entry:
+  %retval = alloca i32, align 4
+  store i32 0, i32* %retval
+  %call = call i32 @_Z3fooi(i32 44)
+  ret i32 %call
+}
diff --git a/llvm/tools/llvm-lto/llvm-lto.cpp b/llvm/tools/llvm-lto/llvm-lto.cpp
index 92bac8351fd..32b31343411 100644
--- a/llvm/tools/llvm-lto/llvm-lto.cpp
+++ b/llvm/tools/llvm-lto/llvm-lto.cpp
@@ -69,6 +69,10 @@ static cl::opt<bool> ListSymbolsOnly(
     "list-symbols-only", cl::init(false),
     cl::desc("Instead of running LTO, list the symbols in each IR file"));
 
+static cl::opt<bool> SetMergedModule(
+    "set-merged-module", cl::init(false),
+    cl::desc("Use the first input module as the merged module"));
+
 namespace {
 struct ModuleInfo {
   std::vector<bool> CanBeHidden;
@@ -194,15 +198,22 @@ int main(int argc, char **argv) {
       return 1;
     }
 
-    if (!CodeGen.addModule(Module.get()))
+    LTOModule *LTOMod = Module.get();
+
+    // We use the first input module as the destination module when
+    // SetMergedModule is true.
+    if (SetMergedModule && i == BaseArg) {
+      // Transfer ownership to the code generator.
+      CodeGen.setModule(Module.release());
+    } else if (!CodeGen.addModule(Module.get()))
       return 1;
 
-    unsigned NumSyms = Module->getSymbolCount();
+    unsigned NumSyms = LTOMod->getSymbolCount();
     for (unsigned I = 0; I < NumSyms; ++I) {
-      StringRef Name = Module->getSymbolName(I);
+      StringRef Name = LTOMod->getSymbolName(I);
       if (!DSOSymbolsSet.count(Name))
         continue;
-      lto_symbol_attributes Attrs = Module->getSymbolAttributes(I);
+      lto_symbol_attributes Attrs = LTOMod->getSymbolAttributes(I);
       unsigned Scope = Attrs & LTO_SYMBOL_SCOPE_MASK;
       if (Scope != LTO_SYMBOL_SCOPE_DEFAULT_CAN_BE_HIDDEN)
         KeptDSOSyms.push_back(Name);
diff --git a/llvm/tools/lto/lto.cpp b/llvm/tools/lto/lto.cpp
index 3a336526e6b..714b8e02d20 100644
--- a/llvm/tools/lto/lto.cpp
+++ b/llvm/tools/lto/lto.cpp
@@ -248,6 +248,10 @@ bool lto_codegen_add_module(lto_code_gen_t cg, lto_module_t mod) {
   return !unwrap(cg)->addModule(unwrap(mod));
 }
 
+void lto_codegen_set_module(lto_code_gen_t cg, lto_module_t mod) {
+  unwrap(cg)->setModule(unwrap(mod));
+}
+
 bool lto_codegen_set_debug_model(lto_code_gen_t cg, lto_debug_model debug) {
   unwrap(cg)->setDebugInfo(debug);
   return false;
diff --git a/llvm/tools/lto/lto.exports b/llvm/tools/lto/lto.exports
index 9ef60772fb8..8729050123e 100644
--- a/llvm/tools/lto/lto.exports
+++ b/llvm/tools/lto/lto.exports
@@ -25,6 +25,7 @@ lto_module_dispose
 lto_api_version
 lto_codegen_set_diagnostic_handler
 lto_codegen_add_module
+lto_codegen_set_module
 lto_codegen_add_must_preserve_symbol
 lto_codegen_compile
 lto_codegen_create