[ELF] - Implemented --retain-symbols-file option

--retain-symbols-file=filename
Retain only the symbols listed in the file filename, discarding all others. 
filename is simply a flat file, with one symbol name per line. This option 
is especially useful in environments (such as VxWorks) where a large global 
symbol table is accumulated gradually, to conserve run-time memory.

Note: though documentation says "--retain-symbols-file does not discard 
undefined symbols, or symbols needed for relocations.", both bfd and gold 
do that, and this patch too, like testcase show.

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

llvm-svn: 290122

5 files changed, 89 insertions, 6 deletions

diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
index d8e78e0e25c..efeadf259af 100644
--- a/lld/ELF/Config.h
+++ b/lld/ELF/Config.h
@@ -12,6 +12,7 @@
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/ELF.h"
 
 #include <vector>
@@ -33,8 +34,8 @@ enum ELFKind {
 // For --build-id.
 enum class BuildIdKind { None, Fast, Md5, Sha1, Hexstring, Uuid };
 
-// For --discard-{all,locals,none}.
-enum class DiscardPolicy { Default, All, Locals, None };
+// For --discard-{all,locals,none} and --retain-symbols-file.
+enum class DiscardPolicy { Default, All, Locals, RetainFile, None };
 
 // For --strip-{all,debug}.
 enum class StripPolicy { None, All, Debug };
@@ -83,6 +84,7 @@ struct Configuration {
   llvm::StringRef OutputFile;
   llvm::StringRef SoName;
   llvm::StringRef Sysroot;
+  llvm::StringSet<> RetainSymbolsFile;
   std::string RPath;
   std::vector<VersionDefinition> VersionDefinitions;
   std::vector<llvm::StringRef> AuxiliaryList;
diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index 935da5dacc3..a567cce8174 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -478,16 +478,35 @@ static SortSectionPolicy getSortKind(opt::InputArgList &Args) {
   return SortSectionPolicy::Default;
 }
 
+static std::vector<StringRef> getLines(MemoryBufferRef MB) {
+  std::vector<StringRef> Ret;
+  SmallVector<StringRef, 0> Arr;
+  MB.getBuffer().split(Arr, '\n');
+  for (StringRef S : Arr) {
+    S = S.trim();
+    if (!S.empty())
+      Ret.push_back(S);
+  }
+  return Ret;
+}
+
 // Parse the --symbol-ordering-file argument. File has form:
 // symbolName1
 // [...]
 // symbolNameN
 static void parseSymbolOrderingList(MemoryBufferRef MB) {
   unsigned I = 0;
-  SmallVector<StringRef, 0> Arr;
-  MB.getBuffer().split(Arr, '\n');
-  for (StringRef S : Arr)
-    Config->SymbolOrderingFile.insert({S.trim(), I++});
+  for (StringRef S : getLines(MB))
+    Config->SymbolOrderingFile.insert({S, I++});
+}
+
+// Parse the --retain-symbols-file argument. File has form:
+// symbolName1
+// [...]
+// symbolNameN
+static void parseRetainSymbolsList(MemoryBufferRef MB) {
+  for (StringRef S : getLines(MB))
+    Config->RetainSymbolsFile.insert(S);
 }
 
 // Initializes Config members by the command line options.
@@ -636,6 +655,14 @@ void LinkerDriver::readConfigs(opt::InputArgList &Args) {
     if (Optional<MemoryBufferRef> Buffer = readFile(Arg->getValue()))
       parseSymbolOrderingList(*Buffer);
 
+  // If --retain-symbol-file is used, we'll retail only the symbols listed in
+  // the file and discard all others.
+  if (auto *Arg = Args.getLastArg(OPT_retain_symbols_file)) {
+    Config->Discard = DiscardPolicy::RetainFile;
+    if (Optional<MemoryBufferRef> Buffer = readFile(Arg->getValue()))
+      parseRetainSymbolsList(*Buffer);
+  }
+
   for (auto *Arg : Args.filtered(OPT_export_dynamic_symbol))
     Config->VersionScriptGlobals.push_back(
         {Arg->getValue(), /*IsExternCpp*/ false, /*HasWildcard*/ false});
diff --git a/lld/ELF/Options.td b/lld/ELF/Options.td
index 98505e6bac0..d7a699ae7c3 100644
--- a/lld/ELF/Options.td
+++ b/lld/ELF/Options.td
@@ -176,6 +176,9 @@ def rpath: S<"rpath">, HelpText<"Add a DT_RUNPATH to the output">;
 
 def relocatable: F<"relocatable">, HelpText<"Create relocatable object file">;
 
+def retain_symbols_file: J<"retain-symbols-file=">, MetaVarName<"<file>">,
+  HelpText<"Retain only the symbols listed in the file">;
+
 def script: S<"script">, HelpText<"Read linker script">;
 
 def section_start: S<"section-start">, MetaVarName<"<address>">,
@@ -269,6 +272,7 @@ def alias_o_output: Joined<["--"], "output=">, Alias<o>;
 def alias_o_output2 : Separate<["--"], "output">, Alias<o>;
 def alias_pie_pic_executable: F<"pic-executable">, Alias<pie>;
 def alias_relocatable_r: Flag<["-"], "r">, Alias<relocatable>;
+def alias_retain_symbols_file: S<"retain-symbols-file">, Alias<retain_symbols_file>;
 def alias_rpath_R: JoinedOrSeparate<["-"], "R">, Alias<rpath>;
 def alias_rpath_rpath: J<"rpath=">, Alias<rpath>;
 def alias_script_T: JoinedOrSeparate<["-"], "T">, Alias<script>;
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index c1640223357..be3a070d082 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -420,6 +420,12 @@ template <class ELFT> static bool includeInSymtab(const SymbolBody &B) {
   if (!B.isLocal() && !B.symbol()->IsUsedInRegularObj)
     return false;
 
+  // If --retain-symbols-file is given, we'll keep only symbols listed in that
+  // file.
+  if (Config->Discard == DiscardPolicy::RetainFile &&
+      !Config->RetainSymbolsFile.count(B.getName()))
+    return false;
+
   if (auto *D = dyn_cast<DefinedRegular<ELFT>>(&B)) {
     // Always include absolute symbols.
     if (!D->Section)
diff --git a/lld/test/ELF/retain-symbols-file.s b/lld/test/ELF/retain-symbols-file.s
new file mode 100644
index 00000000000..b146bd797b5
--- /dev/null
+++ b/lld/test/ELF/retain-symbols-file.s
@@ -0,0 +1,44 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t
+# RUN: echo "bar" > %t_retain.txt
+# RUN: echo "foo" >> %t_retain.txt
+# RUN: ld.lld -shared --retain-symbols-file=%t_retain.txt %t -o %t2
+# RUN: llvm-readobj -s -sd -t %t2 | FileCheck %s
+
+## Check separate form.
+# RUN: ld.lld -shared --retain-symbols-file %t_retain.txt %t -o %t2
+# RUN: llvm-readobj -s -sd -t %t2 | FileCheck %s
+
+# CHECK:       Symbols [
+# CHECK-NEXT:   Symbol {
+# CHECK-NEXT:     Name:  (0)
+# CHECK:        Symbol {
+# CHECK-NEXT:     Name: bar
+# CHECK:        Symbol {
+# CHECK-NEXT:     Name: foo
+# CHECK-NOT:    Symbol
+
+.text
+.globl _start
+_start:
+call zed@PLT
+call und@PLT
+
+.globl foo
+.type foo,@function
+foo:
+retq
+
+.globl bar
+.type bar,@function
+bar:
+retq
+
+.globl zed
+.type zed,@function
+zed:
+retq
+
+.type loc,@function
+loc:
+retq