lld::COFF: better behavior when using as a library

[ruiu]

Summary

Previously, the COFF driver would call exit(0) when called
as a library. Now it takes ExitEarly option, and if it
is false, it doesn't exit. So it is now more library-friendly.

Furthermore, link() calls freeArena() before returning, to
clean up resources.

Based on an Andrew Kelley's patch.

[andrewrk]

lld/COFF/Driver.cpp

64

is this necessary?

[ruiu]

lld/COFF/Driver.cpp

64

Yes, because we use it in Errors.cpp.

lld/COFF/Config.h

struct Configuration {
   uint64_t StackReserve = 1024 * 1024;
   uint64_t StackCommit = 4096;
   uint64_t HeapReserve = 1024 * 1024;
   uint64_t HeapCommit = 4096;
   uint32_t MajorImageVersion = 0;
   uint32_t MinorImageVersion = 0;
   uint32_t MajorOSVersion = 6;
   uint32_t MinorOSVersion = 0;
+  bool CanExitEarly = false;
   bool DynamicBase = true;
   bool AllowBind = true;
   bool NxCompat = true;
   bool AllowIsolation = true;
   bool TerminalServerAware = true;
   bool LargeAddressAware = false;
   bool HighEntropyVA = false;
   bool AppContainer = false;

lld/COFF/Driver.cpp

 
 Configuration *Config;
 LinkerDriver *Driver;
 
 BumpPtrAllocator BAlloc;
 StringSaver Saver{BAlloc};
 std::vector<SpecificAllocBase *> SpecificAllocBase::Instances;
 
-bool link(ArrayRef<const char *> Args, raw_ostream &Diag) {
+bool link(ArrayRef<const char *> Args, bool CanExitEarly, raw_ostream &Diag) {
   ErrorCount = 0;
   ErrorOS = &Diag;
 
   Config = make<Configuration>();
   Config->Argv = {Args.begin(), Args.end()};
   Config->ColorDiagnostics = ErrorOS->has_colors();
+  Config->CanExitEarly = CanExitEarly;
 
   Symtab = make<SymbolTable>();
 
   Driver = make<LinkerDriver>();
   Driver->link(Args);
+
+  // Call exit() if we can to avoid calling destructors.
+  if (CanExitEarly)
+    exitLld(ErrorCount ? 1 : 0);
+
+  freeArena();
   return !ErrorCount;
 }
 
 // Drop directory components and replace extension with ".exe" or ".dll".
 static std::string getOutputPath(StringRef Path) {
   auto P = Path.find_last_of("\\/");
   StringRef S = (P == StringRef::npos) ? Path : Path.substr(P + 1);
   const char* E = Config->DLL ? ".dll" : ".exe";
if (auto *Arg = Args.getLastArg(OPT_deffile)) {
     // parseModuleDefs mutates Config object.
     parseModuleDefs(Arg->getValue());
   }
 
   // Handle generation of import library from a def file.
   if (!Args.hasArg(OPT_INPUT)) {
     fixupExports();
     createImportLibrary(/*AsLib=*/true);
-    exit(0);
+    return;
   }
 
   // Handle /delayload
   for (auto *Arg : Args.filtered(OPT_delayload)) {
     Config->DelayLoads.insert(StringRef(Arg->getValue()).lower());
     if (Config->Machine == I386) {
       Config->DelayLoadHelper = addUndefined("___delayLoadHelper2@8");
     } else {
void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
 
   if (ErrorCount)
     return;
 
   // If /msvclto is given, we use the MSVC linker to link LTO output files.
   // This is useful because MSVC link.exe can generate complete PDBs.
   if (Args.hasArg(OPT_msvclto)) {
     invokeMSVC(Args);
-    exit(0);
+    return;
   }
 
   // Do LTO by compiling bitcode input files to a set of native COFF files then
   // link those files.
   Symtab->addCombinedLTOObjects();
   run();
 
   // Make sure we have resolved all symbols.
if (Config->DoGC)
     markLive(Symtab->getChunks());
 
   // Identify identical COMDAT sections to merge them.
   if (Config->DoICF)
     doICF(Symtab->getChunks());
 
   // Write the result.
   writeResult();
-
-  if (ErrorCount)
-    return;
-
-  // Call exit to avoid calling destructors.
-  exit(0);
 }
 
 } // namespace coff
 } // namespace lld

lld/COFF/Error.h

 void log(const Twine &Msg);
 void message(const Twine &Msg);
 void warn(const Twine &Msg);
 void error(const Twine &Msg);
 LLVM_ATTRIBUTE_NORETURN void fatal(const Twine &Msg);
 LLVM_ATTRIBUTE_NORETURN void fatal(std::error_code EC, const Twine &Prefix);
 LLVM_ATTRIBUTE_NORETURN void fatal(llvm::Error &Err, const Twine &Prefix);
 
+LLVM_ATTRIBUTE_NORETURN void exitLld(int Val);
+
 template <class T> T check(ErrorOr<T> V, const Twine &Prefix) {
   if (auto EC = V.getError())
     fatal(EC, Prefix);
   return std::move(*V);
 }
 
 template <class T> T check(Expected<T> E, const Twine &Prefix) {
   if (llvm::Error Err = E.takeError())

lld/COFF/Error.cpp

 // The functions defined in this file can be called from multiple threads,
 // but outs() or errs() are not thread-safe. We protect them using a mutex.
 static std::mutex Mu;
 
 namespace coff {
 uint64_t ErrorCount;
 raw_ostream *ErrorOS;
 
-static LLVM_ATTRIBUTE_NORETURN void exitLld(int Val) {
+LLVM_ATTRIBUTE_NORETURN void exitLld(int Val) {
   // Dealloc/destroy ManagedStatic variables before calling
   // _exit(). In a non-LTO build, this is a nop. In an LTO
   // build allows us to get the output of -time-passes.
   llvm_shutdown();
 
   outs().flush();
   errs().flush();
   _exit(Val);
void error(const Twine &Msg) {
 
   if (Config->ErrorLimit == 0 || ErrorCount < Config->ErrorLimit) {
     print("error: ", raw_ostream::RED);
     *ErrorOS << Msg << "\n";
   } else if (ErrorCount == Config->ErrorLimit) {
     print("error: ", raw_ostream::RED);
     *ErrorOS << "too many errors emitted, stopping now"
              << " (use /ERRORLIMIT:0 to see all errors)\n";
+    if (Config->CanExitEarly)
       exitLld(1);
   }
 
   ++ErrorCount;
 }
 
 void fatal(const Twine &Msg) {
   if (Config->ColorDiagnostics) {
     errs().changeColor(raw_ostream::RED, /*bold=*/true);

lld/MinGW/Driver.cpp

bool mingw::link(ArrayRef<const char *> ArgsArr, raw_ostream &Diag) {
 
   if (Args.hasArg(OPT__HASH_HASH_HASH))
     return true;
 
   // Repack vector of strings to vector of const char pointers for coff::link.
   std::vector<const char *> Vec;
   for (const std::string &S : LinkArgs)
     Vec.push_back(S.c_str());
-  return coff::link(Vec);
+  return coff::link(Vec, true);
 }

lld/include/lld/Common/Driver.h

 #ifndef LLD_COMMON_DRIVER_H
 #define LLD_COMMON_DRIVER_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace lld {
 namespace coff {
-bool link(llvm::ArrayRef<const char *> Args,
+bool link(llvm::ArrayRef<const char *> Args, bool CanExitEarly,
           llvm::raw_ostream &Diag = llvm::errs());
 }
 
 namespace mingw {
 bool link(llvm::ArrayRef<const char *> Args,
           llvm::raw_ostream &Diag = llvm::errs());
 }
 

lld/tools/lld/lld.cpp

int main(int Argc, const char **Argv) {
 
   std::vector<const char *> Args(Argv, Argv + Argc);
   switch (parseFlavor(Args)) {
   case Gnu:
     if (isPETarget(Args))
       return !mingw::link(Args);
     return !elf::link(Args, true);
   case WinLink:
-    return !coff::link(Args);
+    return !coff::link(Args, true);
   case Darwin:
     return !mach_o::link(Args);
   default:
     die("lld is a generic driver.\n"
         "Invoke ld.lld (Unix), ld (macOS) or lld-link (Windows) instead.");
   }
 }