[WebAssembly] Add support for --whole-archive.

Subscribers: dschuff, jgravelle-google, aheejin, sunfish, llvm-commits

Differential Revision: https://reviews.llvm.org/D49706

5 files changed, 105 insertions, 14 deletions

diff --git a/lld/test/wasm/whole-archive.test b/lld/test/wasm/whole-archive.test
new file mode 100644
index 00000000000..814acbf432e
--- /dev/null
+++ b/lld/test/wasm/whole-archive.test
@@ -0,0 +1,34 @@
+RUN: llc -filetype=obj %p/Inputs/start.ll -o %t.o
+RUN: llc -filetype=obj %p/Inputs/ret32.ll -o %t.ret32.o
+RUN: rm -f %t.a
+RUN: llvm-ar rcs %t.a %t.ret32.o
+
+Should not add symbols from the archive by default as they are not required
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o %t.a
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=NOTADDED %s
+NOTADDED: FunctionNames:
+NOTADDED-NOT: Name: ret32
+NOTADDED: ...
+
+Should add symbols from the archive if --whole-archive is used
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o --whole-archive %t.a
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=ADDED %s
+ADDED: FunctionNames:
+ADDED:   Name: ret32
+ADDED: ...
+
+--no-whole-archive should restore default behaviour
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o --whole-archive --no-whole-archive %t.a
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=NOTADDED %s
+
+--whole-archive and --no-whole-archive should affect only archives which follow them
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o %t.a --whole-archive --no-whole-archive
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=NOTADDED %s
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o --whole-archive %t.a --no-whole-archive
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=ADDED %s
+
+--whole-archive should also work with thin archives
+RUN: rm -f %tthin.a
+RUN: llvm-ar --format=gnu rcsT %tthin.a %t.ret32.o
+RUN: wasm-ld --no-gc-sections -o %t.wasm %t.o --whole-archive %tthin.a
+RUN: obj2yaml %t.wasm | FileCheck --check-prefix=ADDED %s
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index ccaf90fc785..329b5ae80a9 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -68,6 +68,10 @@ private:
   void createFiles(opt::InputArgList &Args);
   void addFile(StringRef Path);
   void addLibrary(StringRef Name);
+
+  // True if we are in --whole-archive and --no-whole-archive.
+  bool InWholeArchive = false;
+
   std::vector<InputFile *> Files;
 };
 } // anonymous namespace
@@ -180,6 +184,37 @@ static void readImportFile(StringRef Filename) {
       Config->AllowUndefinedSymbols.insert(Sym);
 }
 
+// Returns slices of MB by parsing MB as an archive file.
+// Each slice consists of a member file in the archive.
+std::vector<MemoryBufferRef> static getArchiveMembers(
+    MemoryBufferRef MB) {
+  std::unique_ptr<Archive> File =
+      CHECK(Archive::create(MB),
+            MB.getBufferIdentifier() + ": failed to parse archive");
+
+  std::vector<MemoryBufferRef> V;
+  Error Err = Error::success();
+  for (const ErrorOr<Archive::Child> &COrErr : File->children(Err)) {
+    Archive::Child C =
+        CHECK(COrErr, MB.getBufferIdentifier() +
+                          ": could not get the child of the archive");
+    MemoryBufferRef MBRef =
+        CHECK(C.getMemoryBufferRef(),
+              MB.getBufferIdentifier() +
+                  ": could not get the buffer for a child of the archive");
+    V.push_back(MBRef);
+  }
+  if (Err)
+    fatal(MB.getBufferIdentifier() + ": Archive::children failed: " +
+          toString(std::move(Err)));
+
+  // Take ownership of memory buffers created for members of thin archives.
+  for (std::unique_ptr<MemoryBuffer> &MB : File->takeThinBuffers())
+    make<std::unique_ptr<MemoryBuffer>>(std::move(MB));
+
+  return V;
+}
+
 void LinkerDriver::addFile(StringRef Path) {
   Optional<MemoryBufferRef> Buffer = readFile(Path);
   if (!Buffer.hasValue())
@@ -188,6 +223,13 @@ void LinkerDriver::addFile(StringRef Path) {
 
   switch (identify_magic(MBRef.getBuffer())) {
   case file_magic::archive: {
+    // Handle -whole-archive.
+    if (InWholeArchive) {
+      for (MemoryBufferRef &M : getArchiveMembers(MBRef))
+        Files.push_back(createObjectFile(M));
+      return;
+    }
+
     SmallString<128> ImportFile = Path;
     path::replace_extension(ImportFile, ".imports");
     if (fs::exists(ImportFile))
@@ -197,10 +239,11 @@ void LinkerDriver::addFile(StringRef Path) {
     return;
   }
   case file_magic::bitcode:
-    Files.push_back(make<BitcodeFile>(MBRef));
+  case file_magic::wasm_object:
+    Files.push_back(createObjectFile(MBRef));
     break;
   default:
-    Files.push_back(make<ObjFile>(MBRef));
+    error("unknown file type: " + MBRef.getBufferIdentifier());
   }
 }
 
@@ -225,6 +268,12 @@ void LinkerDriver::createFiles(opt::InputArgList &Args) {
     case OPT_INPUT:
       addFile(Arg->getValue());
       break;
+    case OPT_whole_archive:
+      InWholeArchive = true;
+      break;
+    case OPT_no_whole_archive:
+      InWholeArchive = false;
+      break;
     }
   }
 }
diff --git a/lld/wasm/InputFiles.cpp b/lld/wasm/InputFiles.cpp
index 0396f299749..53a24c3cffd 100644
--- a/lld/wasm/InputFiles.cpp
+++ b/lld/wasm/InputFiles.cpp
@@ -42,6 +42,17 @@ Optional<MemoryBufferRef> lld::wasm::readFile(StringRef Path) {
   return MBRef;
 }
 
+InputFile *lld::wasm::createObjectFile(MemoryBufferRef MB) {
+  file_magic Magic = identify_magic(MB.getBuffer());
+  if (Magic == file_magic::wasm_object)
+    return make<ObjFile>(MB);
+
+  if (Magic == file_magic::bitcode)
+    return make<BitcodeFile>(MB);
+
+  fatal("unknown file type: " + MB.getBufferIdentifier());
+}
+
 void ObjFile::dumpInfo() const {
   log("info for: " + getName() +
       "\n              Symbols : " + Twine(Symbols.size()) +
@@ -360,18 +371,7 @@ void ArchiveFile::addMember(const Archive::Symbol *Sym) {
             "could not get the buffer for the member defining symbol " +
                 Sym->getName());
 
-  InputFile *Obj;
-
-  file_magic Magic = identify_magic(MB.getBuffer());
-  if (Magic == file_magic::wasm_object) {
-    Obj = make<ObjFile>(MB);
-  } else if (Magic == file_magic::bitcode) {
-    Obj = make<BitcodeFile>(MB);
-  } else {
-    error("unknown file type: " + MB.getBufferIdentifier());
-    return;
-  }
-
+  InputFile *Obj = createObjectFile(MB);
   Obj->ArchiveName = getName();
   Symtab->addFile(Obj);
 }
diff --git a/lld/wasm/InputFiles.h b/lld/wasm/InputFiles.h
index 73cc32f8add..ec77446e630 100644
--- a/lld/wasm/InputFiles.h
+++ b/lld/wasm/InputFiles.h
@@ -150,6 +150,10 @@ public:
   std::unique_ptr<llvm::lto::InputFile> Obj;
 };
 
+// Will report a fatal() error if the input buffer is not a valid bitcode
+// or was object file.
+InputFile *createObjectFile(MemoryBufferRef MB);
+
 // Opens a given file.
 llvm::Optional<MemoryBufferRef> readFile(StringRef Path);
 
diff --git a/lld/wasm/Options.td b/lld/wasm/Options.td
index d627d7a11d4..43588a830e3 100644
--- a/lld/wasm/Options.td
+++ b/lld/wasm/Options.td
@@ -132,6 +132,10 @@ def no_entry: F<"no-entry">,
 def stack_first: F<"stack-first">,
   HelpText<"Place stack at start of linear memory rather than after data">;
 
+defm whole_archive: B<"whole-archive",
+    "Force load of all members in a static library",
+    "Do not force load of all members in a static library (default)">;
+
 // Aliases
 def alias_entry_e: JoinedOrSeparate<["-"], "e">, Alias<entry>;
 def alias_entry_entry: J<"entry=">, Alias<entry>;