Linker: Add flag to override linkage rules

Add a flag to lib/Linker (and `llvm-link`) to override linkage rules.
When set, the functions in the source module *always* replace those in
the destination module.

The `llvm-link` option is `-override=abc.ll`.  All the "regular" modules
are loaded and linked first, followed by the `-override` modules.  This
is useful for debugging workflows where some subset of the module (e.g.,
a single function) is extracted into a separate file where it's
optimized differently, before being merged back in.

Patch by Luqman Aden!

llvm-svn: 235473

include/llvm/Linker/Linker.h

@@ -68,8 +68,10 @@ public:
   void deleteModule();
 
   /// \brief Link \p Src into the composite. The source is destroyed.
+  /// Passing OverrideSymbols as true will have symbols from Src
+  /// shadow those in the Dest.
   /// Returns true on error.
-  bool linkInModule(Module *Src);
+  bool linkInModule(Module *Src, bool OverrideSymbols = false);
 
   /// \brief Set the composite to the passed-in module.
   void setModule(Module *Dst);

lib/Linker/LinkModules.cpp

@@ -424,12 +424,17 @@ class ModuleLinker {
 
   DiagnosticHandlerFunction DiagnosticHandler;
 
+  /// For symbol clashes, prefer those from Src.
+  bool OverrideFromSrc;
+
 public:
   ModuleLinker(Module *dstM, Linker::IdentifiedStructTypeSet &Set, Module *srcM,
-               DiagnosticHandlerFunction DiagnosticHandler)
+               DiagnosticHandlerFunction DiagnosticHandler,
+               bool OverrideFromSrc)
       : DstM(dstM), SrcM(srcM), TypeMap(Set),
         ValMaterializer(TypeMap, DstM, LazilyLinkGlobalValues),
-        DiagnosticHandler(DiagnosticHandler) {}
+        DiagnosticHandler(DiagnosticHandler), OverrideFromSrc(OverrideFromSrc) {
+  }
 
   bool run();
 
@@ -725,6 +730,12 @@ bool ModuleLinker::getComdatResult(const Comdat *SrcC,
 bool ModuleLinker::shouldLinkFromSource(bool &LinkFromSrc,
                                         const GlobalValue &Dest,
                                         const GlobalValue &Src) {
+  // Should we unconditionally use the Src?
+  if (OverrideFromSrc) {
+    LinkFromSrc = true;
+    return false;
+  }
+
   // We always have to add Src if it has appending linkage.
   if (Src.hasAppendingLinkage()) {
     LinkFromSrc = true;
@@ -1071,8 +1082,9 @@ bool ModuleLinker::linkGlobalValueProto(GlobalValue *SGV) {
   } else {
     // If the GV is to be lazily linked, don't create it just yet.
     // The ValueMaterializerTy will deal with creating it if it's used.
-    if (!DGV && (SGV->hasLocalLinkage() || SGV->hasLinkOnceLinkage() ||
-                 SGV->hasAvailableExternallyLinkage())) {
+    if (!DGV && !OverrideFromSrc &&
+        (SGV->hasLocalLinkage() || SGV->hasLinkOnceLinkage() ||
+         SGV->hasAvailableExternallyLinkage())) {
       DoNotLinkFromSource.insert(SGV);
       return false;
     }
@@ -1738,9 +1750,9 @@ void Linker::deleteModule() {
   Composite = nullptr;
 }
 
-bool Linker::linkInModule(Module *Src) {
+bool Linker::linkInModule(Module *Src, bool OverrideSymbols) {
   ModuleLinker TheLinker(Composite, IdentifiedStructTypes, Src,
-                         DiagnosticHandler);
+                         DiagnosticHandler, OverrideSymbols);
   bool RetCode = TheLinker.run();
   Composite->dropTriviallyDeadConstantArrays();
   return RetCode;

test/Linker/Inputs/override-different-linkage.ll

@@ -0,0 +1,4 @@
+define linkonce i32 @foo(i32 %i) {
+entry:
+  ret i32 4
+}

test/Linker/Inputs/override-with-internal-linkage-2.ll

@@ -0,0 +1,4 @@
+define internal i32 @foo(i32 %i) {
+entry:
+  ret i32 4
+}

test/Linker/Inputs/override-with-internal-linkage.ll

@@ -0,0 +1,4 @@
+define i32 @foo(i32 %i) {
+entry:
+  ret i32 4
+}

test/Linker/Inputs/override.ll

@@ -0,0 +1,4 @@
+define i32 @foo(i32 %i) {
+entry:
+  ret i32 4
+}

test/Linker/override-different-linkage.ll

@@ -0,0 +1,19 @@
+; RUN: llvm-link %s -override %S/Inputs/override-different-linkage.ll -S | FileCheck %s
+; RUN: llvm-link -override %S/Inputs/override-different-linkage.ll %s -S | FileCheck %s
+
+
+; CHECK-LABEL: define linkonce i32 @foo
+; CHECK-NEXT: entry:
+; CHECK-NEXT: ret i32 4
+define weak i32 @foo(i32 %i) {
+entry:
+  %add = add nsw i32 %i, %i
+  ret i32 %add
+}
+
+; Function Attrs: nounwind ssp uwtable
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %a = call i32 @foo(i32 2)
+  ret i32 %a
+}

test/Linker/override-with-internal-linkage-2.ll

@@ -0,0 +1,23 @@
+; RUN: llvm-link %s -override %S/Inputs/override-with-internal-linkage-2.ll -S | FileCheck %s
+; RUN: llvm-link -override %S/Inputs/override-with-internal-linkage-2.ll %s -S | FileCheck %s
+
+; CHECK-LABEL: define i32 @foo
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %add = add nsw i32 %i, %i
+; CHECK-NEXT: ret i32 %add
+define i32 @foo(i32 %i) {
+entry:
+  %add = add nsw i32 %i, %i
+  ret i32 %add
+}
+
+; CHECK-LABEL: define internal i32 @foo1
+; CHECK-NEXT: entry:
+; CHECK-NEXT: ret i32 4
+
+; Function Attrs: nounwind ssp uwtable
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %a = call i32 @foo(i32 2)
+  ret i32 %a
+}

test/Linker/override-with-internal-linkage.ll

@@ -0,0 +1,23 @@
+; RUN: llvm-link %s -override %S/Inputs/override-with-internal-linkage.ll -S | FileCheck %s
+; RUN: llvm-link -override %S/Inputs/override-with-internal-linkage.ll %s -S | FileCheck %s
+
+; CHECK-LABEL: define internal i32 @foo2
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %add = add nsw i32 %i, %i
+; CHECK-NEXT: ret i32 %add
+
+; CHECK-LABEL: define i32 @foo
+; CHECK-NEXT: entry:
+; CHECK-NEXT: ret i32 4
+define internal i32 @foo(i32 %i) {
+entry:
+  %add = add nsw i32 %i, %i
+  ret i32 %add
+}
+
+; Function Attrs: nounwind ssp uwtable
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %a = call i32 @foo(i32 2)
+  ret i32 %a
+}

test/Linker/override.ll

@@ -0,0 +1,19 @@
+; RUN: llvm-link %s -override %S/Inputs/override.ll -S | FileCheck %s
+; RUN: llvm-link -override %S/Inputs/override.ll %s -S | FileCheck %s
+
+
+; CHECK-LABEL: define i32 @foo
+; CHECK-NEXT: entry:
+; CHECK-NEXT: ret i32 4
+define i32 @foo(i32 %i) {
+entry:
+  %add = add nsw i32 %i, %i
+  ret i32 %add
+}
+
+; Function Attrs: nounwind ssp uwtable
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %a = call i32 @foo(i32 2)
+  ret i32 %a
+}

tools/llvm-link/llvm-link.cpp

@@ -38,6 +38,11 @@ static cl::list<std::string>
 InputFilenames(cl::Positional, cl::OneOrMore,
                cl::desc("<input bitcode files>"));
 
+static cl::list<std::string> OverridingInputs(
+    "override", cl::ZeroOrMore, cl::value_desc("filename"),
+    cl::desc(
+        "input bitcode file which can override previously defined symbol(s)"));
+
 static cl::opt<std::string>
 OutputFilename("o", cl::desc("Override output filename"), cl::init("-"),
                cl::value_desc("filename"));
@@ -108,7 +113,8 @@ static void diagnosticHandler(const DiagnosticInfo &DI) {
 }
 
 static bool linkFiles(const char *argv0, LLVMContext &Context, Linker &L,
-                      const cl::list<std::string> &Files) {
+                      const cl::list<std::string> &Files,
+                      bool OverrideDuplicateSymbols) {
   for (const auto &File : Files) {
     std::unique_ptr<Module> M = loadFile(argv0, File, Context);
     if (!M.get()) {
@@ -124,7 +130,7 @@ static bool linkFiles(const char *argv0, LLVMContext &Context, Linker &L,
     if (Verbose)
       errs() << "Linking in '" << File << "'\n";
 
-    if (L.linkInModule(M.get()))
+    if (L.linkInModule(M.get(), OverrideDuplicateSymbols))
       return false;
   }
 
@@ -143,7 +149,12 @@ int main(int argc, char **argv) {
   auto Composite = make_unique<Module>("llvm-link", Context);
   Linker L(Composite.get(), diagnosticHandler);
 
-  if (!linkFiles(argv[0], Context, L, InputFilenames))
+  // First add all the regular input files
+  if (!linkFiles(argv[0], Context, L, InputFilenames, false))
+    return 1;
+
+  // Next the -override ones.
+  if (!linkFiles(argv[0], Context, L, OverridingInputs, true))
     return 1;
 
   if (DumpAsm) errs() << "Here's the assembly:\n" << *Composite;