Re-land [LLD] Allow usage of LLD as a library

This reverts commit aa495214b3.

As discussed in https://github.com/llvm/llvm-project/issues/53475 this patch
allows for 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 hopefully, in subsequent patches, for COFF at least.

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

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

lld/include/lld/Common/Driver.h

@@ -13,8 +13,25 @@
 #include "llvm/Support/raw_ostream.h"
 
 namespace lld {
-struct SafeReturn {
-  int ret;
+enum Flavor {
+  Invalid,
+  Gnu,     // -flavor gnu
+  MinGW,   // -flavor gnu MinGW
+  WinLink, // -flavor link
+  Darwin,  // -flavor darwin
+  Wasm,    // -flavor wasm
+};
+
+using Driver = bool (*)(llvm::ArrayRef<const char *>, llvm::raw_ostream &,
+                        llvm::raw_ostream &, bool, bool);
+
+struct DriverDef {
+  Flavor f;
+  Driver d;
+};
+
+struct Result {
+  int retCode;
   bool canRunAgain;
 };
 
@@ -24,33 +41,29 @@ struct SafeReturn {
 // and re-entry would not be possible anymore. Use exitLld() in that case to
 // properly exit your application and avoid intermittent crashes on exit caused
 // by cleanup.
-SafeReturn safeLldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
-                       llvm::raw_ostream &stderrOS);
-
-namespace coff {
-bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
-}
-
-namespace mingw {
-bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
-}
-
-namespace elf {
-bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
-}
-
-namespace macho {
-bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
-}
-
-namespace wasm {
-bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
-          llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
-}
+Result lldMain(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
+               llvm::raw_ostream &stderrOS, llvm::ArrayRef<DriverDef> drivers);
 } // namespace lld
 
+// With this macro, library users must specify which drivers they use, provide
+// that information to lldMain() in the `drivers` param, and link the
+// corresponding driver library in their executable.
+#define LLD_HAS_DRIVER(name)                                                   \
+  namespace lld {                                                              \
+  namespace name {                                                             \
+  bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,    \
+            llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);  \
+  }                                                                            \
+  }
+
+// An array which declares that all LLD drivers are linked in your executable.
+// Must be used along with LLD_HAS_DRIVERS. See examples in LLD unittests.
+#define LLD_ALL_DRIVERS                                                        \
+  {                                                                            \
+    {lld::WinLink, &lld::coff::link}, {lld::Gnu, &lld::elf::link},             \
+        {lld::MinGW, &lld::mingw::link}, {lld::Darwin, &lld::macho::link}, {   \
+      lld::Wasm, &lld::wasm::link                                              \
+    }                                                                          \
+  }
+
 #endif

lld/test/CMakeLists.txt

@@ -14,12 +14,23 @@ configure_lit_site_cfg(
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
   MAIN_CONFIG
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
-  )
+)
+configure_lit_site_cfg(
+  ${CMAKE_CURRENT_SOURCE_DIR}/Unit/lit.site.cfg.py.in
+  ${CMAKE_CURRENT_BINARY_DIR}/Unit/lit.site.cfg.py
+  MAIN_CONFIG
+  ${CMAKE_CURRENT_SOURCE_DIR}/Unit/lit.cfg.py
+  PATHS
+  "SHLIBDIR"
+  "LLD_BINARY_DIR"
+  "LLD_SOURCE_DIR"
+)
 
 set(LLD_TEST_DEPS lld)
 if (NOT LLD_BUILT_STANDALONE)
   list(APPEND LLD_TEST_DEPS
     FileCheck
+    LLDUnitTests
     count
     dsymutil
     llc
@@ -60,6 +71,7 @@ add_lit_testsuite(check-lld "Running lld test suite"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${LLD_TEST_DEPS}
   )
+set_target_properties(check-lld PROPERTIES FOLDER "lld tests")
 
 add_custom_target(lld-test-depends DEPENDS ${LLD_TEST_DEPS})
 set_target_properties(lld-test-depends PROPERTIES FOLDER "lld tests")
@@ -68,8 +80,6 @@ add_lit_testsuites(LLD ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${LLD_TEST_DEPS}
   )
 
-set_target_properties(check-lld PROPERTIES FOLDER "lld tests")
-
 # Add a legacy target spelling: lld-test
 add_custom_target(lld-test)
 add_dependencies(lld-test check-lld)

lld/test/Unit/lit.cfg.py

@@ -0,0 +1,47 @@
+# -*- Python -*-
+
+# Configuration file for the 'lit' test runner.
+
+import os
+import subprocess
+
+import lit.formats
+
+# name: The name of this test suite.
+config.name = "LLD-Unit"
+
+# suffixes: A list of file extensions to treat as test files.
+config.suffixes = []
+
+# test_source_root: The root path where tests are located.
+# test_exec_root: The root path where tests should be run.
+config.test_exec_root = os.path.join(config.lld_obj_root, "unittests")
+config.test_source_root = config.test_exec_root
+
+# testFormat: The test format to use to interpret tests.
+config.test_format = lit.formats.GoogleTest(config.llvm_build_mode, "Tests")
+
+# Propagate the temp directory. Windows requires this because it uses \Windows\
+# if none of these are present.
+if "TMP" in os.environ:
+    config.environment["TMP"] = os.environ["TMP"]
+if "TEMP" in os.environ:
+    config.environment["TEMP"] = os.environ["TEMP"]
+
+# Propagate HOME as it can be used to override incorrect homedir in passwd
+# that causes the tests to fail.
+if "HOME" in os.environ:
+    config.environment["HOME"] = os.environ["HOME"]
+
+# Win32 seeks DLLs along %PATH%.
+if sys.platform in ["win32", "cygwin"] and os.path.isdir(config.shlibdir):
+    config.environment["PATH"] = os.path.pathsep.join((
+            config.shlibdir, config.environment["PATH"]))
+
+# Win32 may use %SYSTEMDRIVE% during file system shell operations, so propogate.
+if sys.platform == "win32" and "SYSTEMDRIVE" in os.environ:
+    config.environment["SYSTEMDRIVE"] = os.environ["SYSTEMDRIVE"]
+
+# Expand the LLD source path so that unittests can use associated input files.
+# (see AsLibELF/ROCm.cpp test)
+config.environment["LLD_SRC_DIR"] = config.lld_src_dir

lld/test/Unit/lit.site.cfg.py.in

@@ -0,0 +1,12 @@
+@LIT_SITE_CFG_IN_HEADER@
+
+import sys
+
+config.llvm_build_mode = lit_config.substitute("@LLVM_BUILD_MODE@")
+config.shlibdir = lit_config.substitute(path(r"@SHLIBDIR@"))
+config.lld_obj_root = path(r"@LLD_BINARY_DIR@")
+config.lld_src_dir = path(r"@LLD_SOURCE_DIR@")
+
+# Let the main config do the real work.
+lit_config.load_config(
+    config, os.path.join(config.lld_src_dir, "test/Unit/lit.cfg.py"))

lld/tools/lld/lld.cpp

@@ -25,7 +25,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "lld/Common/CommonLinkerContext.h"
 #include "lld/Common/Driver.h"
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
@@ -49,161 +48,15 @@ using namespace lld;
 using namespace llvm;
 using namespace llvm::sys;
 
-enum Flavor {
-  Invalid,
-  Gnu,     // -flavor gnu
-  WinLink, // -flavor link
-  Darwin,  // -flavor darwin
-  Wasm,    // -flavor wasm
-};
+namespace lld {
+extern bool inTestOutputDisabled;
 
-[[noreturn]] static void die(const Twine &s) {
-  llvm::errs() << s << "\n";
-  exit(1);
-}
-
-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 bool isPETarget(std::vector<const char *> &v) {
-  for (auto it = v.begin(); it + 1 != v.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(v.data(), v.data() + v.size());
-  BumpPtrAllocator a;
-  StringSaver saver(a);
-  cl::ExpansionContext ECtx(saver.getAllocator(), getDefaultQuotingStyle());
-  if (Error Err = ECtx.expandResponseFiles(expandedArgs))
-    die(toString(std::move(Err)));
-  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 parseFlavor(std::vector<const char *> &v) {
-  // Parse -flavor option.
-  if (v.size() > 1 && v[1] == StringRef("-flavor")) {
-    if (v.size() <= 2)
-      die("missing arg value for '-flavor'");
-    Flavor f = getFlavor(v[2]);
-    if (f == Invalid)
-      die("Unknown flavor: " + StringRef(v[2]));
-    v.erase(v.begin() + 1, v.begin() + 3);
-    return f;
-  }
-
-  // Deduct the flavor from argv[0].
-  StringRef arg0 = path::filename(v[0]);
-  if (arg0.ends_with_insensitive(".exe"))
-    arg0 = arg0.drop_back(4);
-  return parseProgname(arg0);
-}
-
-bool inTestOutputDisabled = false;
-
-/// Universal linker main(). This linker emulates the gnu, darwin, or
-/// windows linker based on the argv[0] or -flavor option.
-static int lldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
-                   llvm::raw_ostream &stderrOS, bool exitEarly = true) {
-  std::vector<const char *> args(argv, argv + argc);
-  auto link = [&args]() {
-    Flavor f = parseFlavor(args);
-    if (f == Gnu && isPETarget(args))
-      return mingw::link;
-    else if (f == Gnu)
-      return elf::link;
-    else if (f == WinLink)
-      return coff::link;
-    else if (f == Darwin)
-      return macho::link;
-    else if (f == Wasm)
-      return lld::wasm::link;
-    else
-      die("lld is a generic driver.\n"
-          "Invoke ld.lld (Unix), ld64.lld (macOS), lld-link (Windows), wasm-ld"
-          " (WebAssembly) instead");
-  }();
-  // Run the driver. If an error occurs, false will be returned.
-  bool r = link(args, stdoutOS, stderrOS, exitEarly, inTestOutputDisabled);
-
-  // Call exit() if we can to avoid calling destructors.
-  if (exitEarly)
-    exitLld(!r ? 1 : 0);
-
-  // Delete the global context and clear the global context pointer, so that it
-  // cannot be accessed anymore.
-  CommonLinkerContext::destroy();
-
-  return !r ? 1 : 0;
-}
-
-// Similar to lldMain except that exceptions are caught.
-SafeReturn lld::safeLldMain(int argc, const char **argv,
-                            llvm::raw_ostream &stdoutOS,
-                            llvm::raw_ostream &stderrOS) {
-  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 = lldMain(argc, argv, stdoutOS, stderrOS, /*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};
-}
+// Bypass the crash recovery handler, which is only meant to be used in
+// LLD-as-lib scenarios.
+int unsafeLldMain(llvm::ArrayRef<const char *> args,
+                  llvm::raw_ostream &stdoutOS, llvm::raw_ostream &stderrOS,
+                  llvm::ArrayRef<DriverDef> drivers, bool exitEarly);
+} // namespace lld
 
 // When in lit tests, tells how many times the LLD tool should re-execute the
 // main loop with the same inputs. When not in test, returns a value of 0 which
@@ -215,6 +68,12 @@ static unsigned inTestVerbosity() {
   return v;
 }
 
+LLD_HAS_DRIVER(coff)
+LLD_HAS_DRIVER(elf)
+LLD_HAS_DRIVER(mingw)
+LLD_HAS_DRIVER(macho)
+LLD_HAS_DRIVER(wasm)
+
 int lld_main(int argc, char **argv, const llvm::ToolContext &) {
   InitLLVM x(argc, argv);
   sys::Process::UseANSIEscapeCodes(true);
@@ -225,11 +84,16 @@ int lld_main(int argc, char **argv, const llvm::ToolContext &) {
     LLVM_BUILTIN_TRAP;
   }
 
+  ArrayRef<const char *> args(argv, argv + argc);
+
   // Not running in lit tests, just take the shortest codepath with global
   // exception handling and no memory cleanup on exit.
-  if (!inTestVerbosity())
-    return lldMain(argc, const_cast<const char **>(argv), llvm::outs(),
-                   llvm::errs());
+  if (!inTestVerbosity()) {
+    int r =
+        lld::unsafeLldMain(args, llvm::outs(), llvm::errs(), LLD_ALL_DRIVERS,
+                           /*exitEarly=*/true);
+    return r;
+  }
 
   std::optional<int> mainRet;
   CrashRecoveryContext::Enable();
@@ -239,16 +103,15 @@ int lld_main(int argc, char **argv, const llvm::ToolContext &) {
     inTestOutputDisabled = (i != 1);
 
     // Execute one iteration.
-    auto r = safeLldMain(argc, const_cast<const char **>(argv), llvm::outs(),
-                         llvm::errs());
+    auto r = lldMain(args, llvm::outs(), llvm::errs(), LLD_ALL_DRIVERS);
     if (!r.canRunAgain)
-      exitLld(r.ret); // Exit now, can't re-execute again.
+      exitLld(r.retCode); // Exit now, can't re-execute again.
 
     if (!mainRet) {
-      mainRet = r.ret;
-    } else if (r.ret != *mainRet) {
+      mainRet = r.retCode;
+    } else if (r.retCode != *mainRet) {
       // Exit now, to fail the tests if the result is different between runs.
-      return r.ret;
+      return r.retCode;
     }
   }
   return *mainRet;

lld/unittests/AsLibAll/AllDrivers.cpp

@@ -0,0 +1,35 @@
+//===- AllDrivers.cpp -------------------------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+// This test shows a typical case where all LLD drivers are linked into the
+// application binary. This is very similar to how lld.exe binary is linked,
+// except that here we cleanup the internal LLD memory context after each call.
+//===----------------------------------------------------------------------===//
+
+#include "lld/Common/Driver.h"
+#include "gmock/gmock.h"
+
+LLD_HAS_DRIVER(coff)
+LLD_HAS_DRIVER(elf)
+LLD_HAS_DRIVER(mingw)
+LLD_HAS_DRIVER(macho)
+LLD_HAS_DRIVER(wasm)
+
+static bool lldInvoke(std::vector<const char *> args) {
+  args.push_back("--version");
+  lld::Result r =
+      lld::lldMain(args, llvm::outs(), llvm::errs(), LLD_ALL_DRIVERS);
+  return !r.retCode && r.canRunAgain;
+}
+
+TEST(AsLib, AllDrivers) {
+  EXPECT_TRUE(lldInvoke({"ld.lld"}));
+  EXPECT_TRUE(lldInvoke({"ld64.lld"}));
+  EXPECT_TRUE(lldInvoke({"ld", "-m", "i386pe"})); // MinGW
+  EXPECT_TRUE(lldInvoke({"lld-link"}));
+  EXPECT_TRUE(lldInvoke({"wasm-ld"}));
+}

lld/unittests/AsLibAll/CMakeLists.txt

@@ -0,0 +1,17 @@
+# Test usage of LLD as a library.
+# This test covers a typical case where all the LLD drivers are pulled into the
+# application executable.
+
+add_lld_unittests(LLDAsLibAllTests
+  AllDrivers.cpp
+)
+
+target_link_libraries(LLDAsLibAllTests
+  PRIVATE
+  lldCommon
+  lldCOFF
+  lldELF
+  lldMachO
+  lldMinGW
+  lldWasm
+)

lld/unittests/AsLibELF/CMakeLists.txt

@@ -0,0 +1,14 @@
+# Test usage of LLD as a library.
+# This test covers a more singular case where only one LLD driver is used in the
+# target application executable.
+
+add_lld_unittests(LLDAsLibELFTests
+  ROCm.cpp
+  SomeDrivers.cpp
+)
+
+target_link_libraries(LLDAsLibELFTests
+  PRIVATE
+  lldCommon
+  lldELF
+)

lld/unittests/AsLibELF/ROCm.cpp

@@ -0,0 +1,74 @@
+//===- ROCm.cpp -------------------------------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+// The purpose of this test is to showcase a more singular usage of LLD as a
+// library, where only one LLD driver is being used (and linked in the target
+// application). We also expect that linking twice the same object files
+// would yield a successfull result. When used as library, LLD always cleans its
+// internal memory context after each linker call.
+//===----------------------------------------------------------------------===//
+
+// When this flag is on, we actually need the MinGW driver library, not the
+// ELF one. Here our test only covers the case where the ELF driver is linked
+// into the unit test binary.
+#ifndef LLD_DEFAULT_LD_LLD_IS_MINGW
+
+#include "lld/Common/Driver.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/Path.h"
+#include "gmock/gmock.h"
+#include <algorithm>
+
+static std::string expand(const char *path) {
+  if (!llvm::StringRef(path).contains("%"))
+    return std::string(path);
+
+  llvm::SmallString<256> thisPath;
+  thisPath.append(getenv("LLD_SRC_DIR"));
+  llvm::sys::path::append(thisPath, "unittests", "AsLibELF");
+
+  std::string expanded(path);
+  expanded.replace(expanded.find("%S"), 2, thisPath.data(), thisPath.size());
+  return expanded;
+}
+
+LLD_HAS_DRIVER(elf)
+
+static bool lldInvoke(const char *inPath, const char *outPath) {
+  std::vector<const char *> args{"ld.lld", "-shared", inPath, "-o", outPath};
+  lld::Result s = lld::lldMain(args, llvm::outs(), llvm::errs(),
+                               {{lld::Gnu, &lld::elf::link}});
+  return !s.retCode && s.canRunAgain;
+}
+
+static bool runLinker(const char *path) {
+  // Create a temp file for HSA code object.
+  int tempHsacoFD = -1;
+  llvm::SmallString<128> tempHsacoFilename;
+  if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFD,
+                                         tempHsacoFilename)) {
+    return false;
+  }
+  llvm::FileRemover cleanupHsaco(tempHsacoFilename);
+  // Invoke lld. Expect a true return value from lld.
+  std::string expandedPath = expand(path);
+  if (!lldInvoke(expandedPath.data(), tempHsacoFilename.c_str())) {
+    llvm::errs() << "Failed to link: " << expandedPath << "\n";
+    return false;
+  }
+  return true;
+}
+
+TEST(AsLib, ROCm) {
+  EXPECT_TRUE(runLinker("%S/Inputs/kernel1.o"));
+  EXPECT_TRUE(runLinker("%S/Inputs/kernel2.o"));
+  EXPECT_TRUE(runLinker("%S/Inputs/kernel1.o"));
+  EXPECT_TRUE(runLinker("%S/Inputs/kernel2.o"));
+}
+#endif

lld/unittests/AsLibELF/SomeDrivers.cpp

@@ -0,0 +1,37 @@
+//===- SomeDrivers.cpp ------------------------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+// In this test we showcase the fact that only one LLD driver can be invoked -
+// the ELF driver that was linked in the test binary. Calling other drivers
+// would return a failure. When using LLD as a library, any driver can be
+// linked into your application.
+//===----------------------------------------------------------------------===//
+
+#include "lld/Common/Driver.h"
+#include "gmock/gmock.h"
+
+LLD_HAS_DRIVER(elf)
+
+static bool lldInvoke(const char *lldExe) {
+  std::vector<const char *> args{lldExe, "--version"};
+  lld::Result s = lld::lldMain(args, llvm::outs(), llvm::errs(),
+                               {{lld::Gnu, &lld::elf::link}});
+  return !s.retCode && s.canRunAgain;
+}
+
+TEST(AsLib, SomeDrivers) {
+  // When this flag is on, we actually need the MinGW driver library, not the
+  // ELF one. Here our test only covers the case where the ELF driver is linked
+  // into the unit test binary.
+#ifndef LLD_DEFAULT_LD_LLD_IS_MINGW
+  EXPECT_TRUE(lldInvoke("ld.lld")); // ELF
+#endif
+  // These drivers are not linked in this unit test.
+  EXPECT_FALSE(lldInvoke("ld64.lld")); // Mach-O
+  EXPECT_FALSE(lldInvoke("lld-link")); // COFF
+  EXPECT_FALSE(lldInvoke("wasm-ld"));  // Wasm
+}

lld/unittests/CMakeLists.txt

@@ -0,0 +1,8 @@
+set_target_properties(LLDUnitTests PROPERTIES FOLDER "lld tests")
+
+function(add_lld_unittests test_dirname)
+  add_unittest(LLDUnitTests ${test_dirname} ${ARGN})
+endfunction()
+
+add_subdirectory(AsLibAll)
+add_subdirectory(AsLibELF)

lld/wasm/Driver.cpp

@@ -84,7 +84,7 @@ private:
 
 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 CommonLinkerContext;
 
   ctx->e.initialize(stdoutOS, stderrOS, exitEarly, disableOutput);

llvm/cmake/modules/AddLLVM.cmake

@@ -2469,3 +2469,13 @@ function(setup_host_tool tool_name setting_name exe_var_name target_var_name)
     add_custom_target(${target_var_name} DEPENDS ${exe_name})
   endif()
 endfunction()
+
+# Adds the unittests folder if gtest is available.
+function(llvm_add_unittests tests_added)
+  if (EXISTS ${LLVM_THIRD_PARTY_DIR}/unittest/googletest/include/gtest/gtest.h)
+    add_subdirectory(unittests)
+    set(${tests_added} ON PARENT_SCOPE)
+  else()
+    message(WARNING "gtest not found, unittests will not be available")
+  endif()
+endfunction()