[lld-link] implement -thinlto-{prefix,object-suffix}-replace

[inglorion]

Summary

Adds the following two options to lld-link:

-thinlto-prefix-replace: allows replacing a prefix in paths generated
for ThinLTO. This can be used to ensure index files and native object
files are stored in unique directories, allowing multiple distributed
ThinLTO links to proceed concurrently.

-thinlto-object-suffix-replace: allows replacing a suffix in object
file paths involved in ThinLTO. This allows minimized index files to
be used for the thin link while storing the paths to the full bitcode
files for subsequent steps (code generation and final linking).

[ruiu]

LGTM

[inglorion]

rebased and updated variable naming convention

lld/COFF/Config.h

struct Configuration {
   llvm::StringMap<int> order;
 
   // Used for /lldmap.
   std::string mapFile;
 
   // Used for /thinlto-index-only:
   llvm::StringRef thinLTOIndexOnlyArg;
 
+  // Used for /thinlto-object-prefix-replace:
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOPrefixReplace;
+
+  // Used for /thinlto-object-suffix-replace:
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOObjectSuffixReplace;
+
   uint64_t imageBase = -1;
   uint64_t fileAlign = 512;
   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;

lld/COFF/Driver.cpp

bool link(ArrayRef<const char *> args, bool canExitEarly, raw_ostream &diag) {
   freeArena();
   ObjFile::instances.clear();
   ImportFile::instances.clear();
   BitcodeFile::instances.clear();
   memset(MergeChunk::instances, 0, sizeof(MergeChunk::instances));
   return !errorCount();
 }
 
+// Parse options of the form "old;new".
+static std::pair<StringRef, StringRef> getOldNewOptions(opt::InputArgList &args,
+                                                        unsigned id) {
+  auto *arg = args.getLastArg(id);
+  if (!arg)
+    return {"", ""};
+
+  StringRef s = arg->getValue();
+  std::pair<StringRef, StringRef> ret = s.split(';');
+  if (ret.second.empty())
+    error(arg->getSpelling() + " expects 'old;new' format, but got " + s);
+  return ret;
+}
+
 // 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";
   return (s.substr(0, s.rfind('.')) + e).str();
 }
 
if (!config->manifestInput.empty() &&
     fatal("/manifestinput: requires /manifest:embed");
   }
 
   config->thinLTOEmitImportsFiles = args.hasArg(OPT_thinlto_emit_imports_files);
   config->thinLTOIndexOnly = args.hasArg(OPT_thinlto_index_only) ||
                              args.hasArg(OPT_thinlto_index_only_arg);
   config->thinLTOIndexOnlyArg =
       args.getLastArgValue(OPT_thinlto_index_only_arg);
+  config->thinLTOPrefixReplace =
+      getOldNewOptions(args, OPT_thinlto_prefix_replace);
+  config->thinLTOObjectSuffixReplace =
+      getOldNewOptions(args, OPT_thinlto_object_suffix_replace);
   // Handle miscellaneous boolean flags.
   config->allowBind = args.hasFlag(OPT_allowbind, OPT_allowbind_no, true);
   config->allowIsolation =
       args.hasFlag(OPT_allowisolation, OPT_allowisolation_no, true);
   config->incremental =
       args.hasFlag(OPT_incremental, OPT_incremental_no,
                    !config->doGC && !config->doICF && !args.hasArg(OPT_order) &&
                        !args.hasArg(OPT_profile));

lld/COFF/InputFiles.h

public:
   static std::vector<BitcodeFile *> instances;
   std::unique_ptr<llvm::lto::InputFile> obj;
 
 private:
   void parse() override;
 
   std::vector<Symbol *> symbols;
 };
+
+std::string replaceThinLTOSuffix(StringRef path);
 } // namespace coff
 
 std::string toString(const coff::InputFile *file);
 } // namespace lld
 
 #endif

lld/COFF/InputFiles.cpp

if (hdr->getType() == llvm::COFF::IMPORT_CODE)
     thunkSym = symtab->addImportThunk(
         name, cast_or_null<DefinedImportData>(impSym), hdr->Machine);
 }
 
 BitcodeFile::BitcodeFile(MemoryBufferRef mb, StringRef archiveName,
                          uint64_t offsetInArchive)
     : InputFile(BitcodeKind, mb) {
   std::string path = mb.getBufferIdentifier().str();
+  if (config->thinLTOIndexOnly)
+    path = replaceThinLTOSuffix(mb.getBufferIdentifier());
 
   // ThinLTO assumes that all MemoryBufferRefs given to it have a unique
   // name. If two archives define two members with the same name, this
   // causes a collision which result in only one of the objects being taken
   // into consideration at LTO time (which very likely causes undefined
   // symbols later in the link stage). So we append file offset to make
   // filename unique.
   MemoryBufferRef mbref(
MachineTypes BitcodeFile::getMachineType() {
   case Triple::arm:
     return ARMNT;
   case Triple::aarch64:
     return ARM64;
   default:
     return IMAGE_FILE_MACHINE_UNKNOWN;
   }
 }
+
+std::string replaceThinLTOSuffix(StringRef path) {
+  StringRef suffix = config->thinLTOObjectSuffixReplace.first;
+  StringRef repl = config->thinLTOObjectSuffixReplace.second;
+
+  if (path.consume_back(suffix))
+    return (path + repl).str();
+  return path;
+}
 } // namespace coff
 } // namespace lld
 
 // Returns the last element of a path, which is supposed to be a filename.
 static StringRef getBasename(StringRef path) {
   return sys::path::filename(path, sys::path::Style::windows);
 }
 

lld/COFF/LTO.cpp

auto ret =
       llvm::make_unique<raw_fd_ostream>(file, ec, sys::fs::OpenFlags::F_None);
   if (ec) {
     error("cannot open " + file + ": " + ec.message());
     return nullptr;
   }
   return ret;
 }
 
+static std::string getThinLTOOutputFile(StringRef path) {
+  return lto::getThinLTOOutputFile(path,
+                                   config->thinLTOPrefixReplace.first,
+                                   config->thinLTOPrefixReplace.second);
+}
+
 static lto::Config createConfig() {
   lto::Config c;
   c.Options = initTargetOptionsFromCodeGenFlags();
 
   // Always emit a section per function/datum with LTO. LLVM LTO should get most
   // of the benefit of linker GC, but there are still opportunities for ICF.
   c.Options.FunctionSections = true;
   c.Options.DataSections = true;
BitcodeCompiler::BitcodeCompiler() {
   if (!config->thinLTOIndexOnlyArg.empty())
     indexFile = openFile(config->thinLTOIndexOnlyArg);
 
   // Initialize ltoObj.
   lto::ThinBackend backend;
   if (config->thinLTOIndexOnly) {
     auto OnIndexWrite = [&](StringRef S) { thinIndices.erase(S); };
     backend = lto::createWriteIndexesThinBackend(
-        "", "", config->thinLTOEmitImportsFiles, indexFile.get(), OnIndexWrite);
+        config->thinLTOPrefixReplace.first, config->thinLTOPrefixReplace.second,
+        config->thinLTOEmitImportsFiles, indexFile.get(), OnIndexWrite);
   } else if (config->thinLTOJobs != 0) {
     backend = lto::createInProcessThinBackend(config->thinLTOJobs);
   }
 
   ltoObj = llvm::make_unique<lto::LTO>(createConfig(), backend,
                                        config->ltoPartitions);
 }
 
std::vector<StringRef> BitcodeCompiler::compile() {
   checkError(ltoObj->run(
       [&](size_t task) {
         return llvm::make_unique<lto::NativeObjectStream>(
             llvm::make_unique<raw_svector_ostream>(buf[task]));
       },
       cache));
 
   // Emit empty index files for non-indexed files
-  for (StringRef S : thinIndices) {
-    std::string Path(S);
-    openFile(Path + ".thinlto.bc");
+  for (StringRef s : thinIndices) {
+    std::string path = getThinLTOOutputFile(s);
+    openFile(path + ".thinlto.bc");
     if (config->thinLTOEmitImportsFiles)
-      openFile(Path + ".imports");
+      openFile(path + ".imports");
   }
 
   // ThinLTO with index only option is required to generate only the index
   // files. After that, we exit from linker and ThinLTO backend runs in a
   // distributed environment.
   if (config->thinLTOIndexOnly) {
     if (indexFile)
       indexFile->close();

lld/COFF/Options.td

def thinlto_emit_imports_files :
     F<"thinlto-emit-imports-files">,
     HelpText<"Emit .imports files with -thinlto-index-only">;
 def thinlto_index_only :
     F<"thinlto-index-only">,
     HelpText<"Instead of linking, emit ThinLTO index files">;
 def thinlto_index_only_arg : P<
     "thinlto-index-only",
     "-thinlto-index-only and also write native module names to file">;
+def thinlto_object_suffix_replace : P<
+    "thinlto-object-suffix-replace",
+    "'old;new' replace old suffix with new suffix in ThinLTO index">;
+def thinlto_prefix_replace: P<
+    "thinlto-prefix-replace",
+    "'old;new' replace old prefix with new prefix in ThinLTO outputs">;
 def dash_dash_version : Flag<["--"], "version">,
   HelpText<"Print version information">;
 defm threads: B<"threads",
     "Run the linker multi-threaded (default)",
     "Do not run the linker multi-threaded">;
 
 // Flags for debugging
 def lldmap : F<"lldmap">;

lld/test/COFF/thinlto-object-suffix-replace.ll

+; REQUIRES: x86
+
+; Test to make sure the thinlto-object-suffix-replace option is handled
+; correctly.
+
+; Generate bitcode file with summary, as well as a minimized bitcode without
+; the debug metadata for the thin link.
+; RUN: opt -thinlto-bc %s -thin-link-bitcode-file=%t1.thinlink.bc -o %t1.obj
+
+; First perform the thin link on the normal bitcode file, and save the
+; resulting index.
+; RUN: lld-link -thinlto-index-only -entry:main %t1.obj -out:%t3.exe
+; RUN: cp %t1.obj.thinlto.bc %t1.obj.thinlto.bc.orig
+
+; Next perform the thin link on the minimized bitcode file, and compare dump
+; of the resulting index to the above dump to ensure they are identical.
+; RUN: rm -f %t1.obj.thinlto.bc
+; Make sure it isn't inadvertently using the regular bitcode file.
+; RUN: rm -f %t1.obj
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-object-suffix-replace:".thinlink.bc;.obj" \
+; RUN:     %t1.thinlink.bc -out:%t3.exe
+; RUN: diff %t1.obj.thinlto.bc.orig %t1.obj.thinlto.bc
+
+; Ensure lld generates error if suffix replace option not in 'old;new' format.
+; RUN: rm -f %t1.obj.thinlto.bc
+; RUN: not lld-link -entry:main -thinlto-index-only \
+; RUN: -thinlto-object-suffix-replace:"abc:def" %t1.thinlink.bc \
+; RUN: -out:%t3.exe 2>&1 | FileCheck %s --check-prefix=ERR1
+; ERR1: -thinlto-object-suffix-replace: expects 'old;new' format, but got abc:def
+
+; If filename does not end with old suffix, no suffix change should occur,
+; so ".thinlto.bc" will simply be appended to the input file name.
+; RUN: rm -f %t1.thinlink.bc.thinlto.bc
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN: -thinlto-object-suffix-replace:".abc;.obj" %t1.thinlink.bc -out:%t3.exe
+; RUN: ls %t1.thinlink.bc.thinlto.bc
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.24215"
+
+define void @main() {
+entry:
+  ret void
+}
+
+!llvm.dbg.cu = !{}
+
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!llvm.module.flags = !{!1}

lld/test/COFF/thinlto-prefix-replace.ll

+; REQUIRES: x86
+; Check that changing the output path via thinlto-prefix-replace works
+; RUN: mkdir -p %t/oldpath
+; RUN: opt -module-summary %s -o %t/oldpath/t.obj
+
+; Ensure that there is no existing file at the new path, so we properly
+; test the creation of the new file there.
+; RUN: rm -f %t/newpath/t.obj.thinlto.bc
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-prefix-replace:"%t/oldpath/;%t/newpath/" %t/oldpath/t.obj \
+; RUN:     -out:%t/t.exe
+; RUN: ls %t/newpath/t.obj.thinlto.bc
+
+; Ensure that lld errors if prefix replace option is not in 'old;new' format.
+; RUN: rm -f %t/newpath/t.obj.thinlto.bc
+; RUN: not lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-prefix-replace:"abc:def" %t/oldpath/t.obj \
+; RUN:     -out:%t/t.exe 2>&1 | FileCheck --check-prefix=ERR %s
+; ERR: -thinlto-prefix-replace: expects 'old;new' format, but got abc:def
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.24215"
+
+define void @main() {
+  ret void
+}