[LLD] Allow usage of LLD as a library

As discussed in llvm#53475 this patch allows using LLD-as-a-lib. It also lets clients link only the drivers that they want (see unit tests).

This also adds the unit test infra as in the other LLVM projects. Among the test coverage, I've added the original issue from @krzysz00, see: https://github.com/ROCmSoftwarePlatform/D108850-lld-bug-reproduction

Important note: this doesn't allow (yet) linking in parallel. This will come a bit later, in subsequent patches, for COFF at last.

Differential revision: https://reviews.llvm.org/D119049

Author: aganea

lld/CMakeLists.txt

@@ -191,6 +191,8 @@ add_subdirectory(Common)
 add_subdirectory(tools/lld)
 
 if (LLVM_INCLUDE_TESTS)
+  add_custom_target(LLDUnitTests)
+  llvm_add_unittests(LLD_UNITTESTS_ADDED)
   add_subdirectory(test)
 endif()
 

lld/COFF/Driver.cpp

@@ -64,7 +64,7 @@ namespace lld::coff {
 
 bool link(ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
           llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput) {
-  // This driver-specific context will be freed later by lldMain().
+  // This driver-specific context will be freed later by unsafeLldMain().
   auto *ctx = new COFFLinkerContext;
 
   ctx->e.initialize(stdoutOS, stderrOS, exitEarly, disableOutput);

lld/Common/CMakeLists.txt

@@ -23,6 +23,7 @@ set_source_files_properties("${version_inc}"
 add_lld_library(lldCommon
   Args.cpp
   CommonLinkerContext.cpp
+  DriverDispatcher.cpp
   DWARF.cpp
   ErrorHandler.cpp
   Filesystem.cpp

lld/Common/DriverDispatcher.cpp

@@ -0,0 +1,203 @@
+//===- DriverDispatcher.cpp - Support using LLD as a library --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "lld/Common/CommonLinkerContext.h"
+#include "lld/Common/Driver.h"
+#include "lld/Common/ErrorHandler.h"
+#include "lld/Common/Memory.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/CrashRecoveryContext.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/TargetParser/Triple.h"
+#include <cstdlib>
+
+using namespace lld;
+using namespace llvm;
+using namespace llvm::sys;
+
+static void err(const Twine &s) { llvm::errs() << s << "\n"; }
+
+static Flavor getFlavor(StringRef s) {
+  return StringSwitch<Flavor>(s)
+      .CasesLower("ld", "ld.lld", "gnu", Gnu)
+      .CasesLower("wasm", "ld-wasm", Wasm)
+      .CaseLower("link", WinLink)
+      .CasesLower("ld64", "ld64.lld", "darwin", Darwin)
+      .Default(Invalid);
+}
+
+static cl::TokenizerCallback getDefaultQuotingStyle() {
+  if (Triple(sys::getProcessTriple()).getOS() == Triple::Win32)
+    return cl::TokenizeWindowsCommandLine;
+  return cl::TokenizeGNUCommandLine;
+}
+
+static bool isPETargetName(StringRef s) {
+  return s == "i386pe" || s == "i386pep" || s == "thumb2pe" || s == "arm64pe";
+}
+
+static std::optional<bool> isPETarget(llvm::ArrayRef<const char *> args) {
+  for (auto it = args.begin(); it + 1 != args.end(); ++it) {
+    if (StringRef(*it) != "-m")
+      continue;
+    return isPETargetName(*(it + 1));
+  }
+
+  // Expand response files (arguments in the form of @<filename>)
+  // to allow detecting the -m argument from arguments in them.
+  SmallVector<const char *, 256> expandedArgs(args.data(),
+                                              args.data() + args.size());
+  BumpPtrAllocator a;
+  StringSaver saver(a);
+  cl::ExpansionContext ectx(saver.getAllocator(), getDefaultQuotingStyle());
+  if (Error e = ectx.expandResponseFiles(expandedArgs)) {
+    err(toString(std::move(e)));
+    return std::nullopt;
+  }
+
+  for (auto it = expandedArgs.begin(); it + 1 != expandedArgs.end(); ++it) {
+    if (StringRef(*it) != "-m")
+      continue;
+    return isPETargetName(*(it + 1));
+  }
+
+#ifdef LLD_DEFAULT_LD_LLD_IS_MINGW
+  return true;
+#else
+  return false;
+#endif
+}
+
+static Flavor parseProgname(StringRef progname) {
+  // Use GNU driver for "ld" by default.
+  if (progname == "ld")
+    return Gnu;
+
+  // Progname may be something like "lld-gnu". Parse it.
+  SmallVector<StringRef, 3> v;
+  progname.split(v, "-");
+  for (StringRef s : v)
+    if (Flavor f = getFlavor(s))
+      return f;
+  return Invalid;
+}
+
+static Flavor
+parseFlavorWithoutMinGW(llvm::SmallVectorImpl<const char *> &argsV) {
+  // Parse -flavor option.
+  if (argsV.size() > 1 && argsV[1] == StringRef("-flavor")) {
+    if (argsV.size() <= 2) {
+      err("missing arg value for '-flavor'");
+      return Invalid;
+    }
+    Flavor f = getFlavor(argsV[2]);
+    if (f == Invalid) {
+      err("Unknown flavor: " + StringRef(argsV[2]));
+      return Invalid;
+    }
+    argsV.erase(argsV.begin() + 1, argsV.begin() + 3);
+    return f;
+  }
+
+  // Deduct the flavor from argv[0].
+  StringRef arg0 = path::filename(argsV[0]);
+  if (arg0.ends_with_insensitive(".exe"))
+    arg0 = arg0.drop_back(4);
+  Flavor f = parseProgname(arg0);
+  if (f == Invalid) {
+    err("lld is a generic driver.\n"
+        "Invoke ld.lld (Unix), ld64.lld (macOS), lld-link (Windows), wasm-ld"
+        " (WebAssembly) instead");
+    return Invalid;
+  }
+  return f;
+}
+
+static Flavor parseFlavor(llvm::SmallVectorImpl<const char *> &argsV) {
+  Flavor f = parseFlavorWithoutMinGW(argsV);
+  if (f == Gnu) {
+    auto isPE = isPETarget(argsV);
+    if (!isPE)
+      return Invalid;
+    if (*isPE)
+      return MinGW;
+  }
+  return f;
+}
+
+static Driver whichDriver(llvm::SmallVectorImpl<const char *> &argsV,
+                          llvm::ArrayRef<DriverDef> drivers) {
+  Flavor f = parseFlavor(argsV);
+  auto it =
+      llvm::find_if(drivers, [=](auto &driverdef) { return driverdef.f == f; });
+  if (it == drivers.end()) {
+    // Driver is invalid or not available in this build.
+    return [](llvm::ArrayRef<const char *>, llvm::raw_ostream &,
+              llvm::raw_ostream &, bool, bool) { return false; };
+  }
+  return it->d;
+}
+
+namespace lld {
+bool inTestOutputDisabled = false;
+
+/// Universal linker main(). This linker emulates the gnu, darwin, or
+/// windows linker based on the argv[0] or -flavor option.
+int unsafeLldMain(llvm::ArrayRef<const char *> args,
+                  llvm::raw_ostream &stdoutOS, llvm::raw_ostream &stderrOS,
+                  llvm::ArrayRef<DriverDef> drivers, bool exitEarly) {
+  SmallVector<const char *, 256> argsV(args);
+  Driver d = whichDriver(argsV, drivers);
+  // Run the driver. If an error occurs, false will be returned.
+  int r = !d(argsV, stdoutOS, stderrOS, exitEarly, inTestOutputDisabled);
+  // At this point 'r' is either 1 for error, and 0 for no error.
+
+  // Call exit() if we can to avoid calling destructors.
+  if (exitEarly)
+    exitLld(r);
+
+  // Delete the global context and clear the global context pointer, so that it
+  // cannot be accessed anymore.
+  CommonLinkerContext::destroy();
+
+  return r;
+}
+} // namespace lld
+
+Result lld::lldMain(llvm::ArrayRef<const char *> args,
+                    llvm::raw_ostream &stdoutOS, llvm::raw_ostream &stderrOS,
+                    llvm::ArrayRef<DriverDef> drivers) {
+  int r = 0;
+  {
+    // The crash recovery is here only to be able to recover from arbitrary
+    // control flow when fatal() is called (through setjmp/longjmp or
+    // __try/__except).
+    llvm::CrashRecoveryContext crc;
+    if (!crc.RunSafely([&]() {
+          r = unsafeLldMain(args, stdoutOS, stderrOS, drivers,
+                            /*exitEarly=*/false);
+        }))
+      return {crc.RetCode, /*canRunAgain=*/false};
+  }
+
+  // Cleanup memory and reset everything back in pristine condition. This path
+  // is only taken when LLD is in test, or when it is used as a library.
+  llvm::CrashRecoveryContext crc;
+  if (!crc.RunSafely([&]() { CommonLinkerContext::destroy(); })) {
+    // The memory is corrupted beyond any possible recovery.
+    return {r, /*canRunAgain=*/false};
+  }
+  return {r, /*canRunAgain=*/true};
+}

lld/ELF/Driver.cpp

@@ -107,10 +107,11 @@ void Ctx::reset() {
   needsTlsLd.store(false, std::memory_order_relaxed);
 }
 
-bool elf::link(ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-               llvm::raw_ostream &stderrOS, bool exitEarly,
-               bool disableOutput) {
-  // This driver-specific context will be freed later by lldMain().
+namespace lld {
+namespace elf {
+bool link(ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
+          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput) {
+  // This driver-specific context will be freed later by unsafeLldMain().
   auto *ctx = new CommonLinkerContext;
 
   ctx->e.initialize(stdoutOS, stderrOS, exitEarly, disableOutput);
@@ -147,6 +148,8 @@ bool elf::link(ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
 
   return errorCount() == 0;
 }
+} // namespace elf
+} // namespace lld
 
 // Parses a linker -m option.
 static std::tuple<ELFKind, uint16_t, uint8_t> parseEmulation(StringRef emul) {

lld/MachO/Driver.cpp

@@ -1367,9 +1367,10 @@ static void handleExplicitExports() {
   }
 }
 
-bool macho::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
-                 llvm::raw_ostream &stderrOS, bool exitEarly,
-                 bool disableOutput) {
+namespace lld {
+namespace macho {
+bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
+          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput) {
   // This driver-specific context will be freed later by lldMain().
   auto *ctx = new CommonLinkerContext;
 
@@ -1968,3 +1969,5 @@ bool macho::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
 
   return errorCount() == 0;
 }
+} // namespace macho
+} // namespace lld

lld/MinGW/Driver.cpp

@@ -157,11 +157,17 @@ searchLibrary(StringRef name, ArrayRef<StringRef> searchPaths, bool bStatic) {
   return "";
 }
 
+namespace lld {
+namespace coff {
+bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
+          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
+}
+
+namespace mingw {
 // Convert Unix-ish command line arguments to Windows-ish ones and
 // then call coff::link.
-bool mingw::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
-                 llvm::raw_ostream &stderrOS, bool exitEarly,
-                 bool disableOutput) {
+bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
+          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput) {
   auto *ctx = new CommonLinkerContext;
   ctx->e.initialize(stdoutOS, stderrOS, exitEarly, disableOutput);
 
@@ -482,3 +488,5 @@ bool mingw::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
 
   return coff::link(vec, stdoutOS, stderrOS, exitEarly, disableOutput);
 }
+} // namespace mingw
+} // namespace lld

lld/docs/NewLLD.rst

@@ -5,7 +5,7 @@ The ELF Linker as a Library
 ---------------------------
 
 You can embed LLD to your program by linking against it and calling the linker's
-entry point function lld::elf::link.
+entry point function `lld::lldMain`.
 
 The current policy is that it is your responsibility to give trustworthy object
 files. The function is guaranteed to return as long as you do not pass corrupted

lld/docs/index.rst

@@ -36,7 +36,7 @@ Features
   external linkers. All you have to do is to construct object files
   and command line arguments just like you would do to invoke an
   external linker and then call the linker's main function,
-  ``lld::elf::link``, from your code.
+  ``lld::lldMain``, from your code.
 
 - It is small. We are using LLVM libObject library to read from object
   files, so it is not a completely fair comparison, but as of February